Use nalgebra matrices for uncommon paths in consolidate blocks

crates/quantum_info/src/convert_2q_block_matrix.rs

@@ -16,7 +16,7 @@ use pyo3::prelude::*;
 
 use num_complex::Complex64;
 use numpy::PyReadonlyArray2;
-use numpy::nalgebra::{Matrix4, MatrixViewMut4};
+use numpy::nalgebra::{Matrix2, Matrix4, MatrixViewMut4};
 use numpy::ndarray::{Array2, ArrayView2};
 use rustworkx_core::petgraph::stable_graph::NodeIndex;
 
@@ -52,6 +52,16 @@ fn matrix4_from_pyreadonly(array: &PyReadonlyArray2<Complex64>) -> Matrix4<Compl
     )
 }
 
+#[inline]
+fn matrix2_from_pyreadonly(array: &PyReadonlyArray2<Complex64>) -> Matrix2<Complex64> {
+    Matrix2::new(
+        *array.get((0, 0)).unwrap(),
+        *array.get((0, 1)).unwrap(),
+        *array.get((1, 0)).unwrap(),
+        *array.get((1, 1)).unwrap(),
+    )
+}
+
 #[inline]
 pub fn get_matrix_from_inst(inst: &PackedInstruction) -> PyResult<Array2<Complex64>> {
     if let Some(mat) = inst.try_matrix() {
@@ -109,6 +119,34 @@ pub fn get_2q_matrix_from_inst(inst: &PackedInstruction) -> PyResult<Matrix4<Com
     })
 }
 
+#[inline]
+pub fn get_1q_matrix_from_inst(inst: &PackedInstruction) -> PyResult<Matrix2<Complex64>> {
+    if let Some(mat) = inst.try_matrix_as_nalgebra_1q() {
+        return Ok(mat);
+    }
+    if inst.op.try_standard_gate().is_some() {
+        return Err(QiskitError::new_err(
+            "Parameterized gates can't be consolidated",
+        ));
+    }
+    let OperationRef::Gate(gate) = inst.op.view() else {
+        return Err(QiskitError::new_err(
+            "Can't compute matrix of non-unitary op",
+        ));
+    };
+    // If the operation is a custom python gate, we will acquire the gil and use an
+    // Operator. Otherwise, using op.matrix() should work.  A user should not be
+    // able to reach this condition in Rust standalone mode.
+    Python::attach(|py| {
+        let res = QI_OPERATOR
+            .get_bound(py)
+            .call1((gate.instruction.clone_ref(py),))?
+            .getattr(intern!(py, "data"))?
+            .extract()?;
+        Ok(matrix2_from_pyreadonly(&res))
+    })
+}
+
 /// Quaternion-based collect of two parallel runs of 1q gates.
 #[derive(Clone, Debug)]
 struct Separable1q {

crates/transpiler/src/passes/consolidate_blocks.rs

@@ -31,7 +31,9 @@ use qiskit_circuit::interner::Interned;
 use qiskit_circuit::operations::StandardGate;
 use qiskit_circuit::operations::{ArrayType, Operation, Param, UnitaryGate};
 use qiskit_circuit::packed_instruction::PackedOperation;
-use qiskit_quantum_info::convert_2q_block_matrix::{blocks_to_matrix, get_matrix_from_inst};
+use qiskit_quantum_info::convert_2q_block_matrix::{
+    blocks_to_matrix, get_1q_matrix_from_inst, get_2q_matrix_from_inst, get_matrix_from_inst,
+};
 use qiskit_synthesis::linalg::nalgebra_array_view;
 use qiskit_synthesis::two_qubit_decompose::RXXEquivalent;
 use qiskit_synthesis::two_qubit_decompose::{
@@ -289,14 +291,31 @@ fn py_run_consolidate_blocks(
                 phys_qargs.get(dag, inst.qubits),
             ) {
                 all_block_gates.insert(inst_node);
-                let matrix = match get_matrix_from_inst(inst) {
-                    Ok(mat) => mat,
-                    Err(_) => continue,
-                };
-                // TODO: Use Matrix2/ArrayType::OneQ when we're using nalgebra
-                // for consolidation
-                let unitary_gate = UnitaryGate {
-                    array: ArrayType::NDArray(matrix),
+                let num_qubits = inst.op.num_qubits();
+                let unitary_gate = if num_qubits == 1 {
+                    let matrix = match get_1q_matrix_from_inst(inst) {
+                        Ok(mat) => mat,
+                        Err(_) => continue,
+                    };
+                    UnitaryGate {
+                        array: ArrayType::OneQ(matrix),
+                    }
+                } else if num_qubits == 2 {
+                    let matrix = match get_2q_matrix_from_inst(inst) {
+                        Ok(mat) => mat,
+                        Err(_) => continue,
+                    };
+                    UnitaryGate {
+                        array: ArrayType::TwoQ(matrix),
+                    }
+                } else {
+                    let matrix = match get_matrix_from_inst(inst) {
+                        Ok(mat) => mat,
+                        Err(_) => continue,
+                    };
+                    UnitaryGate {
+                        array: ArrayType::NDArray(matrix),
+                    }
                 };
                 dag.substitute_op(
                     inst_node,
@@ -450,12 +469,15 @@ fn py_run_consolidate_blocks(
                     first_qubits,
                 )
             {
-                let matrix = match get_matrix_from_inst(first_inst) {
+                // Runs are necessarily single qubit gates so if it has a matrix then it
+                // must have a single qubit matrix or it's not a run and the blocks handling above
+                // would have covered it
+                let matrix = match get_1q_matrix_from_inst(first_inst) {
                     Ok(mat) => mat,
                     Err(_) => continue,
                 };
                 let unitary_gate = UnitaryGate {
-                    array: ArrayType::NDArray(matrix),
+                    array: ArrayType::OneQ(matrix),
                 };
                 dag.substitute_op(
                     first_inst_node,