Leverage native UnitaryGate from rust

crates/accelerate/Cargo.toml

@@ -31,6 +31,7 @@ ndarray_einsum_beta = "0.7"
 once_cell = "1.20.2"
 rustiq-core = "0.0.10"
 bytemuck.workspace = true
+nalgebra.workspace = true
 
 [dependencies.smallvec]
 workspace = true

crates/accelerate/src/circuit_library/quantum_volume.rs

@@ -15,19 +15,15 @@ use pyo3::prelude::*;
 use pyo3::types::PyDict;
 
 use crate::getenv_use_multiple_threads;
-use faer_ext::{IntoFaerComplex, IntoNdarrayComplex};
-use ndarray::prelude::*;
-use num_complex::Complex64;
-use numpy::IntoPyArray;
+use nalgebra::Matrix4;
+use num_complex::{Complex64, ComplexFloat};
 use rand::prelude::*;
 use rand_distr::StandardNormal;
 use rand_pcg::Pcg64Mcg;
 use rayon::prelude::*;
 
 use qiskit_circuit::circuit_data::CircuitData;
-use qiskit_circuit::imports::UNITARY_GATE;
-use qiskit_circuit::operations::Param;
-use qiskit_circuit::operations::PyGate;
+use qiskit_circuit::operations::{ArrayType, Param, UnitaryGate};
 use qiskit_circuit::packed_instruction::PackedOperation;
 use qiskit_circuit::{Clbit, Qubit};
 use smallvec::{smallvec, SmallVec};
@@ -50,11 +46,11 @@ fn random_complex(rng: &mut Pcg64Mcg) -> Complex64 {
 //
 // https://github.com/scipy/scipy/blob/v1.14.1/scipy/stats/_multivariate.py#L4224-L4256
 #[inline]
-fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Array2<Complex64>> {
+fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Matrix4<Complex64>> {
     let mut rng = Pcg64Mcg::seed_from_u64(seed);
 
     (0..size).map(move |_| {
-        let raw_numbers: [[Complex64; 4]; 4] = [
+        let mat: Matrix4<Complex64> = [
             [
                 random_complex(&mut rng),
                 random_complex(&mut rng),
@@ -79,23 +75,11 @@ fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Array2<Compl
                 random_complex(&mut rng),
                 random_complex(&mut rng),
             ],
-        ];
-
-        let qr = aview2(&raw_numbers).into_faer_complex().qr();
-        let r = qr.compute_r();
-        let diag: [Complex64; 4] = [
-            r[(0, 0)].to_num_complex() / r[(0, 0)].abs(),
-            r[(1, 1)].to_num_complex() / r[(1, 1)].abs(),
-            r[(2, 2)].to_num_complex() / r[(2, 2)].abs(),
-            r[(3, 3)].to_num_complex() / r[(3, 3)].abs(),
-        ];
-        let mut q = qr.compute_q().as_ref().into_ndarray_complex().to_owned();
-        q.axis_iter_mut(Axis(0)).for_each(|mut row| {
-            row.iter_mut()
-                .enumerate()
-                .for_each(|(index, val)| *val *= diag[index])
-        });
-        q
+        ]
+        .into();
+        let (q, r) = mat.qr().unpack();
+        let diag = r.map_diagonal(|x| x / x.abs());
+        q.map_with_location(|i, _j, val| val * diag[i])
     })
 }
 
@@ -115,29 +99,21 @@ pub fn quantum_volume(
 
     let kwargs = PyDict::new(py);
     kwargs.set_item(intern!(py, "num_qubits"), 2)?;
-    let mut build_instruction = |(unitary_index, unitary_array): (usize, Array2<Complex64>),
+    let mut build_instruction = |(unitary_index, unitary_array): (usize, Matrix4<Complex64>),
                                  rng: &mut Pcg64Mcg|
      -> PyResult<Instruction> {
         let layer_index = unitary_index % width;
         if layer_index == 0 {
             permutation.shuffle(rng);
         }
-        let unitary = unitary_array.into_pyarray(py);
 
-        let unitary_gate = UNITARY_GATE
-            .get_bound(py)
-            .call((unitary.clone(), py.None(), false), Some(&kwargs))?;
-        let instruction = PyGate {
-            qubits: 2,
-            clbits: 0,
-            params: 1,
-            op_name: "unitary".to_string(),
-            gate: unitary_gate.unbind(),
+        let unitary_gate = UnitaryGate {
+            array: ArrayType::TwoQ(unitary_array),
         };
         let qubit = layer_index * 2;
         Ok((
-            PackedOperation::from_gate(Box::new(instruction)),
-            smallvec![Param::Obj(unitary.into_any().unbind())],
+            PackedOperation::from_unitary(Box::new(unitary_gate)),
+            smallvec![],
             vec![permutation[qubit], permutation[qubit + 1]],
             vec![],
         ))
@@ -156,7 +132,7 @@ pub fn quantum_volume(
         .take(num_unitaries)
         .collect();
 
-    let unitaries: Vec<Array2<Complex64>> = if getenv_use_multiple_threads() && num_unitaries > 200
+    let unitaries: Vec<Matrix4<Complex64>> = if getenv_use_multiple_threads() && num_unitaries > 200
     {
         seed_vec
             .par_chunks(per_thread)

crates/accelerate/src/consolidate_blocks.rs

@@ -11,18 +11,22 @@
 // that they have been altered from the originals.
 
 use hashbrown::{HashMap, HashSet};
+use nalgebra::Matrix2;
 use ndarray::{aview2, Array2};
 use num_complex::Complex64;
-use numpy::{IntoPyArray, PyReadonlyArray2};
+use numpy::PyReadonlyArray2;
 use pyo3::intern;
 use pyo3::prelude::*;
 use rustworkx_core::petgraph::stable_graph::NodeIndex;
+use smallvec::smallvec;
 
 use qiskit_circuit::circuit_data::CircuitData;
+use qiskit_circuit::circuit_instruction::ExtraInstructionAttributes;
 use qiskit_circuit::dag_circuit::DAGCircuit;
 use qiskit_circuit::gate_matrix::{ONE_QUBIT_IDENTITY, TWO_QUBIT_IDENTITY};
-use qiskit_circuit::imports::{QI_OPERATOR, QUANTUM_CIRCUIT, UNITARY_GATE};
-use qiskit_circuit::operations::{Operation, Param};
+use qiskit_circuit::imports::{QI_OPERATOR, QUANTUM_CIRCUIT};
+use qiskit_circuit::operations::{ArrayType, Operation, Param, UnitaryGate};
+use qiskit_circuit::packed_instruction::PackedOperation;
 use qiskit_circuit::Qubit;
 
 use crate::convert_2q_block_matrix::{blocks_to_matrix, get_matrix_from_inst};
@@ -112,11 +116,17 @@ pub(crate) fn consolidate_blocks(
                     Ok(mat) => mat,
                     Err(_) => continue,
                 };
-                let array = matrix.into_pyarray(py);
-                let unitary_gate = UNITARY_GATE
-                    .get_bound(py)
-                    .call1((array, py.None(), false))?;
-                dag.substitute_node_with_py_op(py, inst_node, &unitary_gate, false)?;
+                // TODO: Use Matrix2/ArrayType::OneQ when we're using nalgebra
+                // for consolidation
+                let unitary_gate = UnitaryGate {
+                    array: ArrayType::NDArray(matrix),
+                };
+                dag.substitute_op(
+                    inst_node,
+                    PackedOperation::from_unitary(Box::new(unitary_gate)),
+                    smallvec![],
+                    ExtraInstructionAttributes::default(),
+                )?;
                 continue;
             }
         }
@@ -180,16 +190,16 @@ pub(crate) fn consolidate_blocks(
                     dag.remove_op_node(node);
                 }
             } else {
-                let unitary_gate = UNITARY_GATE.get_bound(py).call1((
-                    array.as_ref().into_pyobject(py)?,
-                    py.None(),
-                    false,
-                ))?;
+                let matrix = array.as_array().to_owned();
+                let unitary_gate = UnitaryGate {
+                    array: ArrayType::NDArray(matrix),
+                };
                 let clbit_pos_map = HashMap::new();
-                dag.replace_block_with_py_op(
-                    py,
+                dag.replace_block(
                     &block,
-                    unitary_gate,
+                    PackedOperation::from_unitary(Box::new(unitary_gate)),
+                    smallvec![],
+                    ExtraInstructionAttributes::default(),
                     false,
                     &block_index_map,
                     &clbit_pos_map,
@@ -213,21 +223,22 @@ pub(crate) fn consolidate_blocks(
                             dag.remove_op_node(node);
                         }
                     } else {
-                        let array = matrix.into_pyarray(py);
-                        let unitary_gate =
-                            UNITARY_GATE
-                                .get_bound(py)
-                                .call1((array, py.None(), false))?;
+                        // TODO: Use Matrix4/ArrayType::TwoQ when we're using nalgebra
+                        // for consolidation
+                        let unitary_gate = UnitaryGate {
+                            array: ArrayType::NDArray(matrix),
+                        };
                         let qubit_pos_map = block_index_map
                             .into_iter()
                             .enumerate()
                             .map(|(idx, qubit)| (qubit, idx))
                             .collect();
                         let clbit_pos_map = HashMap::new();
-                        dag.replace_block_with_py_op(
-                            py,
+                        dag.replace_block(
                             &block,
-                            unitary_gate,
+                            PackedOperation::from_unitary(Box::new(unitary_gate)),
+                            smallvec![],
+                            ExtraInstructionAttributes::default(),
                             false,
                             &qubit_pos_map,
                             &clbit_pos_map,
@@ -258,11 +269,15 @@ pub(crate) fn consolidate_blocks(
                     Ok(mat) => mat,
                     Err(_) => continue,
                 };
-                let array = matrix.into_pyarray(py);
-                let unitary_gate = UNITARY_GATE
-                    .get_bound(py)
-                    .call1((array, py.None(), false))?;
-                dag.substitute_node_with_py_op(py, first_inst_node, &unitary_gate, false)?;
+                let unitary_gate = UnitaryGate {
+                    array: ArrayType::NDArray(matrix),
+                };
+                dag.substitute_op(
+                    first_inst_node,
+                    PackedOperation::from_unitary(Box::new(unitary_gate)),
+                    smallvec![],
+                    ExtraInstructionAttributes::default(),
+                )?;
                 continue;
             }
             let qubit = first_qubits[0];
@@ -293,17 +308,19 @@ pub(crate) fn consolidate_blocks(
                     dag.remove_op_node(node);
                 }
             } else {
-                let array = aview2(&matrix).to_owned().into_pyarray(py);
-                let unitary_gate = UNITARY_GATE
-                    .get_bound(py)
-                    .call1((array, py.None(), false))?;
+                let array: Matrix2<Complex64> =
+                    Matrix2::from_row_iterator(matrix.into_iter().flat_map(|x| x.into_iter()));
+                let unitary_gate = UnitaryGate {
+                    array: ArrayType::OneQ(array),
+                };
                 let mut block_index_map: HashMap<Qubit, usize> = HashMap::with_capacity(1);
                 block_index_map.insert(qubit, 0);
                 let clbit_pos_map = HashMap::new();
-                dag.replace_block_with_py_op(
-                    py,
+                dag.replace_block(
                     &run,
-                    unitary_gate,
+                    PackedOperation::from_unitary(Box::new(unitary_gate)),
+                    smallvec![],
+                    ExtraInstructionAttributes::default(),
                     false,
                     &block_index_map,
                     &clbit_pos_map,

crates/accelerate/src/split_2q_unitaries.rs

@@ -10,15 +10,15 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
-use pyo3::intern;
+use nalgebra::Matrix2;
+use num_complex::Complex64;
 use pyo3::prelude::*;
-use pyo3::types::PyDict;
 use rustworkx_core::petgraph::stable_graph::NodeIndex;
+use smallvec::{smallvec, SmallVec};
 
-use qiskit_circuit::circuit_instruction::OperationFromPython;
 use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType, Wire};
-use qiskit_circuit::imports::UNITARY_GATE;
-use qiskit_circuit::operations::{Operation, Param};
+use qiskit_circuit::operations::{ArrayType, Operation, OperationRef, Param, UnitaryGate};
+use qiskit_circuit::packed_instruction::PackedOperation;
 
 use crate::two_qubit_decompose::{Specialization, TwoQubitWeylDecomposition};
 
@@ -39,7 +39,7 @@ pub fn split_2q_unitaries(
             // We only attempt to split UnitaryGate objects, but this could be extended in future
             // -- however we need to ensure that we can compile the resulting single-qubit unitaries
             // to the supported basis gate set.
-            if qubits.len() != 2 || inst.op.name() != "unitary" {
+            if qubits.len() != 2 || !matches!(inst.op.view(), OperationRef::Unitary(_)) {
                 continue;
             }
             let matrix = inst
@@ -52,22 +52,33 @@ pub fn split_2q_unitaries(
                 None,
             )?;
             if matches!(decomp.specialization, Specialization::IdEquiv) {
-                let k1r_arr = decomp.K1r(py);
-                let k1l_arr = decomp.K1l(py);
-                let kwargs = PyDict::new(py);
-                kwargs.set_item(intern!(py, "num_qubits"), 1)?;
-                let k1r_gate = UNITARY_GATE
-                    .get_bound(py)
-                    .call((k1r_arr, py.None(), false), Some(&kwargs))?;
-                let k1l_gate = UNITARY_GATE
-                    .get_bound(py)
-                    .call((k1l_arr, py.None(), false), Some(&kwargs))?;
-                let insert_fn = |edge: &Wire| -> PyResult<OperationFromPython> {
+                let k1r_arr = decomp.k1r_view();
+                let k1l_arr = decomp.k1l_view();
+
+                let insert_fn = |edge: &Wire| -> (PackedOperation, SmallVec<[Param; 3]>) {
                     if let Wire::Qubit(qubit) = edge {
                         if *qubit == qubits[0] {
-                            k1r_gate.extract()
+                            let mat: Matrix2<Complex64> = [
+                                [k1r_arr[[0, 0]], k1r_arr[[0, 1]]],
+                                [k1r_arr[[1, 0]], k1r_arr[[1, 1]]],
+                            ]
+                            .into();
+                            let k1r_gate = Box::new(UnitaryGate {
+                                array: ArrayType::OneQ(mat),
+                            });
+                            (PackedOperation::from_unitary(k1r_gate), smallvec![])
                         } else {
-                            k1l_gate.extract()
+                            let mat: Matrix2<Complex64> = [
+                                [k1l_arr[[0, 0]], k1l_arr[[0, 1]]],
+                                [k1l_arr[[1, 0]], k1l_arr[[1, 1]]],
+                            ]
+                            .into();
+
+                            let k1l_gate = Box::new(UnitaryGate {
+                                array: ArrayType::OneQ(mat),
+                            });
+
+                            (PackedOperation::from_unitary(k1l_gate), smallvec![])
                         }
                     } else {
                         unreachable!("This will only be called on ops with no classical wires.");

crates/accelerate/src/two_qubit_decompose.rs

@@ -531,6 +531,23 @@ impl TwoQubitWeylDecomposition {
     pub fn c(&self) -> f64 {
         self.c
     }
+
+    pub fn k1l_view(&self) -> ArrayView2<Complex64> {
+        self.K1l.view()
+    }
+
+    pub fn k2l_view(&self) -> ArrayView2<Complex64> {
+        self.K2l.view()
+    }
+
+    pub fn k1r_view(&self) -> ArrayView2<Complex64> {
+        self.K1r.view()
+    }
+
+    pub fn k2r_view(&self) -> ArrayView2<Complex64> {
+        self.K2r.view()
+    }
+
     fn weyl_gate(
         &self,
         simplify: bool,

crates/accelerate/src/unitary_synthesis.rs

@@ -289,7 +289,7 @@ fn py_run_main_loop(
                 py_op: new_node.unbind().into(),
             };
         }
-        if !(packed_instr.op.name() == "unitary"
+        if !(matches!(packed_instr.op.view(), OperationRef::Unitary(_))
             && packed_instr.op.num_qubits() >= min_qubits as u32)
         {
             out_dag.push_back(py, packed_instr)?;