Optimize ConsolidateBlocks further

crates/accelerate/src/convert_2q_block_matrix.rs

@@ -0,0 +1,103 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2022
+//
+// 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::*;
+use pyo3::wrap_pyfunction;
+use pyo3::Python;
+
+use num_complex::Complex64;
+use numpy::ndarray::linalg::kron;
+use numpy::ndarray::{array, Array, Array2, ArrayView2};
+use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2};
+
+/// Return the matrix Operator resulting from a block of Instructions.
+#[pyfunction]
+#[pyo3(text_signature = "(matrix_list, /")]
+pub fn blocks_to_matrix(
+    py: Python,
+    op_list: Vec<(PyReadonlyArray2<Complex64>, Vec<usize>)>,
+) -> PyResult<Py<PyArray2<Complex64>>> {
+    let mut matrix: Array2<Complex64> = Array::eye(4);
+    let swap_gate = array![
+        [
+            Complex64::new(1., 0.),
+            Complex64::new(0., 0.),
+            Complex64::new(0., 0.),
+            Complex64::new(0., 0.)
+        ],
+        [
+            Complex64::new(0., 0.),
+            Complex64::new(0., 0.),
+            Complex64::new(1., 0.),
+            Complex64::new(0., 0.)
+        ],
+        [
+            Complex64::new(0., 0.),
+            Complex64::new(1., 0.),
+            Complex64::new(0., 0.),
+            Complex64::new(0., 0.)
+        ],
+        [
+            Complex64::new(0., 0.),
+            Complex64::new(0., 0.),
+            Complex64::new(0., 0.),
+            Complex64::new(1., 0.)
+        ]
+    ];
+    let identity: Array2<Complex64> = Array::eye(2);
+    for (op_matrix, q_list) in op_list {
+        let op_matrix = op_matrix.as_array();
+        let result = calculate_matrix(op_matrix, &q_list, &swap_gate, &identity);
+        matrix = match result {
+            Some(result) => result.dot(&matrix),
+            None => op_matrix.dot(&matrix),
+        };
+    }
+    Ok(matrix.into_pyarray(py).to_owned())
+}
+
+/// Performs the matrix operations for an Instruction in a 2 qubit system
+fn calculate_matrix(
+    matrix: ArrayView2<Complex64>,
+    q_list: &Vec<usize>,
+    swap_gate: &Array2<Complex64>,
+    identity: &Array2<Complex64>,
+) -> Option<Array2<Complex64>> {
+    let mut basis_change = false;
+    let current: Option<Array2<Complex64>> = if q_list.len() < 2 {
+        if q_list[0] == 1 {
+            Some(kron(&matrix, identity))
+        } else {
+            Some(kron(identity, &matrix))
+        }
+    } else {
+        if q_list[0] > q_list[1] && matrix != swap_gate {
+            basis_change = true;
+        }
+        None
+    };
+
+    if basis_change {
+        match current {
+            Some(current) => Some(swap_gate.dot(&current).dot(swap_gate)),
+            None => Some(swap_gate.dot(&matrix).dot(swap_gate)),
+        }
+    } else {
+        current
+    }
+}
+
+#[pymodule]
+pub fn convert_2q_block_matrix(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(blocks_to_matrix))?;
+    Ok(())
+}

crates/accelerate/src/lib.rs

@@ -16,6 +16,7 @@ use pyo3::prelude::*;
 use pyo3::wrap_pymodule;
 use pyo3::Python;
 
+mod convert_2q_block_matrix;
 mod dense_layout;
 mod edge_collections;
 mod error_map;
@@ -61,5 +62,8 @@ fn _accelerate(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pymodule!(
         euler_one_qubit_decomposer::euler_one_qubit_decomposer
     ))?;
+    m.add_wrapped(wrap_pymodule!(
+        convert_2q_block_matrix::convert_2q_block_matrix
+    ))?;
     Ok(())
 }

qiskit/__init__.py

@@ -40,6 +40,9 @@
 sys.modules[
     "qiskit._accelerate.euler_one_qubit_decomposer"
 ] = qiskit._accelerate.euler_one_qubit_decomposer
+sys.modules[
+    "qiskit._accelerate.convert_2q_block_matrix"
+] = qiskit._accelerate.convert_2q_block_matrix
 
 
 # Extend namespace for backwards compat

qiskit/transpiler/passes/optimization/consolidate_blocks.py

@@ -11,6 +11,7 @@
 # that they have been altered from the originals.
 
 """Replace each block of consecutive gates by a single Unitary node."""
+from __future__ import annotations
 
 import numpy as np
 

qiskit/transpiler/passes/utils/block_to_matrix.py

@@ -12,15 +12,9 @@
 
 """Converts any block of 2 qubit gates into a matrix."""
 
-from numpy import identity, kron
-from qiskit.circuit.library import SwapGate
 from qiskit.quantum_info import Operator
 from qiskit.exceptions import QiskitError
-
-
-SWAP_GATE = SwapGate()
-SWAP_MATRIX = SWAP_GATE.to_matrix()
-IDENTITY = identity(2, dtype=complex)
+from qiskit._accelerate.convert_2q_block_matrix import blocks_to_matrix
 
 
 def _block_to_matrix(block, block_index_map):
@@ -35,36 +29,19 @@ def _block_to_matrix(block, block_index_map):
     Returns:
         NDArray: Matrix representation of the block of operations.
     """
+    op_list = []
     block_index_length = len(block_index_map)
     if block_index_length != 2:
         raise QiskitError(
             "This function can only operate with blocks of 2 qubits."
             + f"This block had {block_index_length}"
         )
-    matrix = identity(2**block_index_length, dtype=complex)
     for node in block:
         try:
             current = node.op.to_matrix()
         except QiskitError:
             current = Operator(node.op).data
         q_list = [block_index_map[qubit] for qubit in node.qargs]
-        if len(q_list) > 2:
-            raise QiskitError(
-                f"The operation {node.op.name} in this block has "
-                + f"{len(q_list)} qubits, only 2 max allowed."
-            )
-        basis_change = False
-        if len(q_list) < block_index_length:
-            if q_list[0] == 1:
-                current = kron(current, IDENTITY)
-            else:
-                current = kron(IDENTITY, current)
-        else:
-            if q_list[0] > q_list[1]:
-                if node.op != SWAP_GATE:
-                    basis_change = True
-        if basis_change:
-            matrix = (SWAP_MATRIX @ current) @ (SWAP_MATRIX @ matrix)
-        else:
-            matrix = current @ matrix
+        op_list.append((current, q_list))
+    matrix = blocks_to_matrix(op_list)
     return matrix