Leverage native UnitaryGate from rust (#13765)

This commit builds off of the native rust representation of a
UnitaryGate added in #13759 and uses the native representation
everywhere we were using UnitaryGate in rust via python previously:
the quantum_volume() function, consolidate blocks, split2qunitaries,
and unitary synthesis.

One future item is consolidate blocks can be updated to use nalgebra
types internally instead of ndarray as for the 1 and 2q cases we know
the fixed size of the array ahead of time. However the block
consolidation code is built using ndarray currently and later synthesis
code also works in ndarray so there isn't any real benefit yet, and we'd
just add unecessary conversions and allocations. However, once #13649
merges this will change and it would make more sense to add the unitary
gate with a Matrix4. But this can be handled separately after this
merges.

Author: mtreinish

Diff for: Cargo.lock

@@ -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

Diff for: crates/accelerate/Cargo.toml

@@ -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)

Diff for: crates/accelerate/src/circuit_library/quantum_volume.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,

Diff for: crates/accelerate/src/consolidate_blocks.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.");

Diff for: crates/accelerate/src/split_2q_unitaries.rs

@@ -533,6 +533,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,

Diff for: crates/accelerate/src/two_qubit_decompose.rs

@@ -332,7 +332,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)?;