From 0a1cddbc8397cacccc8b241e8c8100894d775bfa Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Mon, 1 Jul 2024 17:41:29 -0400 Subject: [PATCH] Use rust gates for ConsolidateBlocks This commit moves to use rust gates for the ConsolidateBlocks transpiler pass. Instead of generating the unitary matrices for the gates in a 2q block Python side and passing that list to a rust function this commit switches to passing a list of DAGOpNodes to the rust and then generating the matrices inside the rust function directly. This is similar to what was done in #12650 for Optimize1qGatesDecomposition. Besides being faster to get the matrix for standard gates, it also reduces the eager construction of Python gate objects which was a significant source of overhead after #12459. To that end this builds on the thread of work in the two PRs #12692 and #12701 which changed the access patterns for other passes to minimize eager gate object construction. --- .../accelerate/src/convert_2q_block_matrix.rs | 63 ++++++++++++++++--- crates/circuit/src/bit_data.rs | 8 ++- crates/circuit/src/circuit_instruction.rs | 5 +- crates/circuit/src/imports.rs | 1 + crates/circuit/src/lib.rs | 2 +- qiskit/dagcircuit/dagcircuit.py | 7 ++- .../passes/optimization/consolidate_blocks.py | 17 ++--- qiskit/transpiler/passes/utils/__init__.py | 1 - .../passes/utils/block_to_matrix.py | 47 -------------- 9 files changed, 78 insertions(+), 73 deletions(-) delete mode 100644 qiskit/transpiler/passes/utils/block_to_matrix.py diff --git a/crates/accelerate/src/convert_2q_block_matrix.rs b/crates/accelerate/src/convert_2q_block_matrix.rs index 9c179397d641..9bec29eba68a 100644 --- a/crates/accelerate/src/convert_2q_block_matrix.rs +++ b/crates/accelerate/src/convert_2q_block_matrix.rs @@ -10,7 +10,9 @@ // copyright notice, and modified files need to carry a notice indicating // that they have been altered from the originals. +use pyo3::intern; use pyo3::prelude::*; +use pyo3::types::PyDict; use pyo3::wrap_pyfunction; use pyo3::Python; @@ -20,32 +22,77 @@ use numpy::ndarray::{aview2, Array2, ArrayView2}; use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; use smallvec::SmallVec; +use qiskit_circuit::bit_data::BitData; +use qiskit_circuit::circuit_instruction::{operation_type_to_py, CircuitInstruction}; +use qiskit_circuit::dag_node::DAGOpNode; use qiskit_circuit::gate_matrix::ONE_QUBIT_IDENTITY; +use qiskit_circuit::imports::QI_OPERATOR; +use qiskit_circuit::operations::{Operation, OperationType}; + +use crate::QiskitError; + +fn get_matrix_from_inst<'py>( + py: Python<'py>, + inst: &'py CircuitInstruction, +) -> PyResult> { + match inst.operation.matrix(&inst.params) { + Some(mat) => Ok(mat), + None => match inst.operation { + OperationType::Standard(_) => Err(QiskitError::new_err( + "Parameterized gates can't be consolidated", + )), + OperationType::Gate(_) => Ok(QI_OPERATOR + .get_bound(py) + .call1((operation_type_to_py(py, inst)?,))? + .getattr(intern!(py, "data"))? + .extract::>()? + .as_array() + .to_owned()), + _ => unreachable!("Only called for unitary ops"), + }, + } +} /// Return the matrix Operator resulting from a block of Instructions. #[pyfunction] #[pyo3(text_signature = "(op_list, /")] pub fn blocks_to_matrix( py: Python, - op_list: Vec<(PyReadonlyArray2, SmallVec<[u8; 2]>)>, + op_list: Vec>, + block_index_map_dict: &Bound, ) -> PyResult>> { + let mut bit_map: BitData = BitData::new(py, "qargs".to_string()); + + for bit in block_index_map_dict.keys() { + bit_map.add(py, &bit, true)?; + } let identity = aview2(&ONE_QUBIT_IDENTITY); - let input_matrix = op_list[0].0.as_array(); - let mut matrix: Array2 = match op_list[0].1.as_slice() { + let first_node = &op_list[0]; + let input_matrix = get_matrix_from_inst(py, &first_node.instruction)?; + let mut matrix: Array2 = match bit_map + .map_bits(first_node.instruction.qubits.bind(py).iter())? + .map(|x| x as u8) + .collect::>() + .as_slice() + { [0] => kron(&identity, &input_matrix), [1] => kron(&input_matrix, &identity), - [0, 1] => input_matrix.to_owned(), - [1, 0] => change_basis(input_matrix), + [0, 1] => input_matrix, + [1, 0] => change_basis(input_matrix.view()), [] => Array2::eye(4), _ => unreachable!(), }; - for (op_matrix, q_list) in op_list.into_iter().skip(1) { - let op_matrix = op_matrix.as_array(); + for node in op_list.into_iter().skip(1) { + let op_matrix = get_matrix_from_inst(py, &node.instruction)?; + let q_list = bit_map + .map_bits(node.instruction.qubits.bind(py).iter())? + .map(|x| x as u8) + .collect::>(); let result = match q_list.as_slice() { [0] => Some(kron(&identity, &op_matrix)), [1] => Some(kron(&op_matrix, &identity)), - [1, 0] => Some(change_basis(op_matrix)), + [1, 0] => Some(change_basis(op_matrix.view())), [] => Some(Array2::eye(4)), _ => None, }; diff --git a/crates/circuit/src/bit_data.rs b/crates/circuit/src/bit_data.rs index 40540f9df5a4..977d1b34e496 100644 --- a/crates/circuit/src/bit_data.rs +++ b/crates/circuit/src/bit_data.rs @@ -70,7 +70,7 @@ impl PartialEq for BitAsKey { impl Eq for BitAsKey {} #[derive(Clone, Debug)] -pub(crate) struct BitData { +pub struct BitData { /// The public field name (i.e. `qubits` or `clbits`). description: String, /// Registered Python bits. @@ -81,7 +81,7 @@ pub(crate) struct BitData { cached: Py, } -pub(crate) struct BitNotFoundError<'py>(pub(crate) Bound<'py, PyAny>); +pub struct BitNotFoundError<'py>(pub(crate) Bound<'py, PyAny>); impl<'py> From> for PyErr { fn from(error: BitNotFoundError) -> Self { @@ -111,6 +111,10 @@ where self.bits.len() } + pub fn is_empty(&self) -> bool { + self.bits.is_empty() + } + /// Gets a reference to the underlying vector of Python bits. #[inline] pub fn bits(&self) -> &Vec { diff --git a/crates/circuit/src/circuit_instruction.rs b/crates/circuit/src/circuit_instruction.rs index 9e4e8fe53cfb..f5c7b8725db8 100644 --- a/crates/circuit/src/circuit_instruction.rs +++ b/crates/circuit/src/circuit_instruction.rs @@ -822,10 +822,7 @@ impl CircuitInstruction { /// Take a reference to a `CircuitInstruction` and convert the operation /// inside that to a python side object. -pub(crate) fn operation_type_to_py( - py: Python, - circuit_inst: &CircuitInstruction, -) -> PyResult { +pub fn operation_type_to_py(py: Python, circuit_inst: &CircuitInstruction) -> PyResult { let (label, duration, unit, condition) = match &circuit_inst.extra_attrs { None => (None, None, None, None), Some(extra_attrs) => ( diff --git a/crates/circuit/src/imports.rs b/crates/circuit/src/imports.rs index e2a1ca542453..39bd2331b0fb 100644 --- a/crates/circuit/src/imports.rs +++ b/crates/circuit/src/imports.rs @@ -72,6 +72,7 @@ pub static SINGLETON_GATE: ImportOnceCell = pub static SINGLETON_CONTROLLED_GATE: ImportOnceCell = ImportOnceCell::new("qiskit.circuit.singleton", "SingletonControlledGate"); pub static DEEPCOPY: ImportOnceCell = ImportOnceCell::new("copy", "deepcopy"); +pub static QI_OPERATOR: ImportOnceCell = ImportOnceCell::new("qiskit.quantum_info", "Operator"); pub static WARNINGS_WARN: ImportOnceCell = ImportOnceCell::new("warnings", "warn"); diff --git a/crates/circuit/src/lib.rs b/crates/circuit/src/lib.rs index 9fcaa36480cf..44239669afae 100644 --- a/crates/circuit/src/lib.rs +++ b/crates/circuit/src/lib.rs @@ -10,6 +10,7 @@ // copyright notice, and modified files need to carry a notice indicating // that they have been altered from the originals. +pub mod bit_data; pub mod circuit_data; pub mod circuit_instruction; pub mod dag_node; @@ -19,7 +20,6 @@ pub mod operations; pub mod parameter_table; pub mod util; -mod bit_data; mod interner; use pyo3::prelude::*; diff --git a/qiskit/dagcircuit/dagcircuit.py b/qiskit/dagcircuit/dagcircuit.py index 0213d242097f..7cc4fa1cc575 100644 --- a/qiskit/dagcircuit/dagcircuit.py +++ b/qiskit/dagcircuit/dagcircuit.py @@ -54,6 +54,7 @@ from qiskit.dagcircuit.dagnode import DAGNode, DAGOpNode, DAGInNode, DAGOutNode from qiskit.circuit.bit import Bit from qiskit.pulse import Schedule +from qiskit._accelerate.circuit import StandardGate, PyGate BitLocations = namedtuple("BitLocations", ("index", "registers")) # The allowable arguments to :meth:`DAGCircuit.copy_empty_like`'s ``vars_mode``. @@ -2169,10 +2170,10 @@ def collect_2q_runs(self): def filter_fn(node): if isinstance(node, DAGOpNode): return ( - isinstance(node.op, Gate) + isinstance(node._raw_op, (StandardGate, PyGate)) and len(node.qargs) <= 2 - and not getattr(node.op, "condition", None) - and not node.op.is_parameterized() + and not getattr(node, "condition", None) + and not node.is_parameterized() ) else: return None diff --git a/qiskit/transpiler/passes/optimization/consolidate_blocks.py b/qiskit/transpiler/passes/optimization/consolidate_blocks.py index 8dca049aa9ac..d2c2d5326717 100644 --- a/qiskit/transpiler/passes/optimization/consolidate_blocks.py +++ b/qiskit/transpiler/passes/optimization/consolidate_blocks.py @@ -27,9 +27,10 @@ from qiskit.circuit.controlflow import ControlFlowOp from qiskit.transpiler.passmanager import PassManager from qiskit.transpiler.passes.synthesis import unitary_synthesis -from qiskit.transpiler.passes.utils import _block_to_matrix from .collect_1q_runs import Collect1qRuns from .collect_2q_blocks import Collect2qBlocks +from qiskit.circuit.controlflow import CONTROL_FLOW_OP_NAMES +from qiskit._accelerate.convert_2q_block_matrix import blocks_to_matrix class ConsolidateBlocks(TransformationPass): @@ -105,14 +106,14 @@ def run(self, dag): block_cargs = set() for nd in block: block_qargs |= set(nd.qargs) - if isinstance(nd, DAGOpNode) and getattr(nd.op, "condition", None): - block_cargs |= set(getattr(nd.op, "condition", None)[0]) + if isinstance(nd, DAGOpNode) and getattr(nd, "condition", None): + block_cargs |= set(getattr(nd, "condition", None)[0]) all_block_gates.add(nd) block_index_map = self._block_qargs_to_indices(dag, block_qargs) for nd in block: - if nd.op.name == basis_gate_name: + if nd.name == basis_gate_name: basis_count += 1 - if self._check_not_in_basis(dag, nd.op.name, nd.qargs): + if self._check_not_in_basis(dag, nd.name, nd.qargs): outside_basis = True if len(block_qargs) > 2: q = QuantumRegister(len(block_qargs)) @@ -124,7 +125,7 @@ def run(self, dag): qc.append(nd.op, [q[block_index_map[i]] for i in nd.qargs]) unitary = UnitaryGate(Operator(qc), check_input=False) else: - matrix = _block_to_matrix(block, block_index_map) + matrix = blocks_to_matrix(block, block_index_map) unitary = UnitaryGate(matrix, check_input=False) max_2q_depth = 20 # If depth > 20, there will be 1q gates to consolidate. @@ -192,7 +193,9 @@ def _handle_control_flow_ops(self, dag): pass_manager.append(Collect2qBlocks()) pass_manager.append(self) - for node in dag.op_nodes(ControlFlowOp): + for node in dag.op_nodes(): + if node.name not in CONTROL_FLOW_OP_NAMES: + continue node.op = node.op.replace_blocks(pass_manager.run(block) for block in node.op.blocks) return dag diff --git a/qiskit/transpiler/passes/utils/__init__.py b/qiskit/transpiler/passes/utils/__init__.py index 7227409a213c..9cba3b1a50cd 100644 --- a/qiskit/transpiler/passes/utils/__init__.py +++ b/qiskit/transpiler/passes/utils/__init__.py @@ -31,4 +31,3 @@ # Utility functions from . import control_flow -from .block_to_matrix import _block_to_matrix diff --git a/qiskit/transpiler/passes/utils/block_to_matrix.py b/qiskit/transpiler/passes/utils/block_to_matrix.py deleted file mode 100644 index 4b8a09cd30f7..000000000000 --- a/qiskit/transpiler/passes/utils/block_to_matrix.py +++ /dev/null @@ -1,47 +0,0 @@ -# This code is part of Qiskit. -# -# (C) Copyright IBM 2017, 2018. -# -# 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. - -"""Converts any block of 2 qubit gates into a matrix.""" - -from qiskit.quantum_info import Operator -from qiskit.exceptions import QiskitError -from qiskit._accelerate.convert_2q_block_matrix import blocks_to_matrix - - -def _block_to_matrix(block, block_index_map): - """ - The function converts any sequence of operations between two qubits into a matrix - that can be utilized to create a gate or a unitary. - - Args: - block (List(DAGOpNode)): A block of operations on two qubits. - block_index_map (dict(Qubit, int)): The mapping of the qubit indices in the main circuit. - - 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}" - ) - 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] - op_list.append((current, q_list)) - matrix = blocks_to_matrix(op_list) - return matrix