From 4fe9dbc6a80dfb113eea6bc18da8bc5e80c4d4f2 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Mon, 8 Jul 2024 15:47:23 -0400 Subject: [PATCH] Use rust gates for ConsolidateBlocks (#12704) * 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. * Add rust filter function for DAGCircuit.collect_2q_runs() * Update crates/accelerate/src/convert_2q_block_matrix.rs --------- Co-authored-by: John Lapeyre --- .../accelerate/src/convert_2q_block_matrix.rs | 104 ++++++++++++++++-- crates/circuit/src/bit_data.rs | 8 +- crates/circuit/src/imports.rs | 1 + crates/circuit/src/lib.rs | 2 +- qiskit/dagcircuit/dagcircuit.py | 24 +--- .../passes/optimization/consolidate_blocks.py | 19 ++-- qiskit/transpiler/passes/utils/__init__.py | 1 - .../passes/utils/block_to_matrix.py | 47 -------- 8 files changed, 120 insertions(+), 86 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..7a9165777dc9 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,84 @@ 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>> { + // Build a BitData in block_index_map_dict order. block_index_map_dict is a dict of bits to + // indices mapping the order of the qargs in the block. There should only be 2 entries since + // there are only 2 qargs here (e.g. `{Qubit(): 0, Qubit(): 1}`) so we need to ensure that + // we added the qubits to bit data in the correct index order. + let mut index_map: Vec = (0..block_index_map_dict.len()).map(|_| py.None()).collect(); + for bit_tuple in block_index_map_dict.items() { + let (bit, index): (PyObject, usize) = bit_tuple.extract()?; + index_map[index] = bit; + } + let mut bit_map: BitData = BitData::new(py, "qargs".to_string()); + for bit in index_map { + bit_map.add(py, bit.bind(py), 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())? + .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, }; @@ -71,8 +125,42 @@ pub fn change_basis(matrix: ArrayView2) -> Array2 { trans_matrix } +#[pyfunction] +pub fn collect_2q_blocks_filter(node: &Bound) -> Option { + match node.downcast::() { + Ok(bound_node) => { + let node = bound_node.borrow(); + match &node.instruction.operation { + OperationType::Standard(gate) => Some( + gate.num_qubits() <= 2 + && node + .instruction + .extra_attrs + .as_ref() + .and_then(|attrs| attrs.condition.as_ref()) + .is_none() + && !node.is_parameterized(), + ), + OperationType::Gate(gate) => Some( + gate.num_qubits() <= 2 + && node + .instruction + .extra_attrs + .as_ref() + .and_then(|attrs| attrs.condition.as_ref()) + .is_none() + && !node.is_parameterized(), + ), + _ => Some(false), + } + } + Err(_) => None, + } +} + #[pymodule] pub fn convert_2q_block_matrix(m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(blocks_to_matrix))?; + m.add_wrapped(wrap_pyfunction!(collect_2q_blocks_filter))?; Ok(()) } 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/imports.rs b/crates/circuit/src/imports.rs index dfefb0a348f8..d277dea9f89c 100644 --- a/crates/circuit/src/imports.rs +++ b/crates/circuit/src/imports.rs @@ -74,6 +74,7 @@ pub static SINGLETON_CONTROLLED_GATE: ImportOnceCell = pub static CONTROLLED_GATE: ImportOnceCell = ImportOnceCell::new("qiskit.circuit", "ControlledGate"); 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"); /// A mapping from the enum variant in crate::operations::StandardGate to the python diff --git a/crates/circuit/src/lib.rs b/crates/circuit/src/lib.rs index 9f0a8017bf21..c7469434c668 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; @@ -20,7 +21,6 @@ pub mod parameter_table; pub mod slice; pub mod util; -mod bit_data; mod interner; use pyo3::prelude::*; diff --git a/qiskit/dagcircuit/dagcircuit.py b/qiskit/dagcircuit/dagcircuit.py index b93a90e47f7b..5d6739d72198 100644 --- a/qiskit/dagcircuit/dagcircuit.py +++ b/qiskit/dagcircuit/dagcircuit.py @@ -55,6 +55,7 @@ from qiskit.circuit.bit import Bit from qiskit.pulse import Schedule from qiskit._accelerate.euler_one_qubit_decomposer import collect_1q_runs_filter +from qiskit._accelerate.convert_2q_block_matrix import collect_2q_blocks_filter BitLocations = namedtuple("BitLocations", ("index", "registers")) # The allowable arguments to :meth:`DAGCircuit.copy_empty_like`'s ``vars_mode``. @@ -1348,9 +1349,9 @@ def replace_block_with_op( for nd in node_block: block_qargs |= set(nd.qargs) block_cargs |= set(nd.cargs) - if (condition := getattr(nd.op, "condition", None)) is not None: + if (condition := getattr(nd, "condition", None)) is not None: block_cargs.update(condition_resources(condition).clbits) - elif isinstance(nd.op, SwitchCaseOp): + elif nd.name in CONTROL_FLOW_OP_NAMES and isinstance(nd.op, SwitchCaseOp): if isinstance(nd.op.target, Clbit): block_cargs.add(nd.op.target) elif isinstance(nd.op.target, ClassicalRegister): @@ -2158,28 +2159,13 @@ def collect_1q_runs(self) -> list[list[DAGOpNode]]: def collect_2q_runs(self): """Return a set of non-conditional runs of 2q "op" nodes.""" - to_qid = {} - for i, qubit in enumerate(self.qubits): - to_qid[qubit] = i - - def filter_fn(node): - if isinstance(node, DAGOpNode): - return ( - isinstance(node.op, Gate) - and len(node.qargs) <= 2 - and not getattr(node.op, "condition", None) - and not node.op.is_parameterized() - ) - else: - return None - def color_fn(edge): if isinstance(edge, Qubit): - return to_qid[edge] + return self.find_bit(edge).index else: return None - return rx.collect_bicolor_runs(self._multi_graph, filter_fn, color_fn) + return rx.collect_bicolor_runs(self._multi_graph, collect_2q_blocks_filter, color_fn) def nodes_on_wire(self, wire, only_ops=False): """ diff --git a/qiskit/transpiler/passes/optimization/consolidate_blocks.py b/qiskit/transpiler/passes/optimization/consolidate_blocks.py index 8dca049aa9ac..72a08efe0f7d 100644 --- a/qiskit/transpiler/passes/optimization/consolidate_blocks.py +++ b/qiskit/transpiler/passes/optimization/consolidate_blocks.py @@ -24,10 +24,11 @@ from qiskit.circuit.library.generalized_gates.unitary import UnitaryGate from qiskit.circuit.library.standard_gates import CXGate from qiskit.transpiler.basepasses import TransformationPass -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 qiskit.circuit.controlflow import CONTROL_FLOW_OP_NAMES +from qiskit._accelerate.convert_2q_block_matrix import blocks_to_matrix + from .collect_1q_runs import Collect1qRuns from .collect_2q_blocks import Collect2qBlocks @@ -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