From a0c3026bc8ad011b85220138c74968cfdc5431f4 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Thu, 26 Sep 2024 07:49:56 -0400 Subject: [PATCH] Fully port BarrierBeforeFinalMeasurements to rust (#13220) * Fully port BarrierBeforeFinalMeasurements to rust This commit migrates the BarrierBeforeFinalMeasurements transpiler pass to operate fully in Rust. The full path of the transpiler pass now never leaves Rust until it has finished modifying the DAGCircuit. The one exception is when the `label` is not set then we still call `MergeAdjacentBarriers` in the python code of the pass. This is the first step in the improvement of the performance of this pass. We can easily leverage multhreading to potentially parallelize the analysis portion of this pass that searches for the final operations and returns the set of indices. But this is blocked on #13219 which prevents us from accessing the PackedInstructions stored in the DAGCircuit in a multithreaded context. This commit also fixes an issue related to shared references in the disjoint_utils module around barrier labels. The combine_barriers() function was incorrectly mutating the label by reference which wouldn't persist in the DAG, and this was causing failures after the barrier was originally generated in Rust with this pass now. Fixes #12253 * Remove unused imports --- .../src/barrier_before_final_measurement.rs | 118 ++++++++++++++++++ crates/accelerate/src/lib.rs | 1 + crates/circuit/src/dag_circuit.rs | 2 +- crates/pyext/src/lib.rs | 1 + qiskit/__init__.py | 3 + .../passes/layout/disjoint_utils.py | 28 +++-- .../barrier_before_final_measurements.py | 58 +-------- 7 files changed, 141 insertions(+), 70 deletions(-) create mode 100644 crates/accelerate/src/barrier_before_final_measurement.rs diff --git a/crates/accelerate/src/barrier_before_final_measurement.rs b/crates/accelerate/src/barrier_before_final_measurement.rs new file mode 100644 index 000000000000..f7143d910b98 --- /dev/null +++ b/crates/accelerate/src/barrier_before_final_measurement.rs @@ -0,0 +1,118 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2024 +// +// 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 hashbrown::HashSet; +use pyo3::prelude::*; +use rustworkx_core::petgraph::stable_graph::NodeIndex; + +use qiskit_circuit::circuit_instruction::ExtraInstructionAttributes; +use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType}; +use qiskit_circuit::imports::BARRIER; +use qiskit_circuit::operations::{Operation, PyInstruction}; +use qiskit_circuit::packed_instruction::{PackedInstruction, PackedOperation}; +use qiskit_circuit::Qubit; + +static FINAL_OP_NAMES: [&str; 2] = ["measure", "barrier"]; + +#[pyfunction] +pub fn barrier_before_final_measurements( + py: Python, + dag: &mut DAGCircuit, + label: Option, +) -> PyResult<()> { + let final_ops: HashSet = dag + .op_nodes(true) + .filter(|node| { + let NodeType::Operation(ref inst) = dag.dag()[*node] else { + unreachable!(); + }; + if !FINAL_OP_NAMES.contains(&inst.op.name()) { + return false; + } + let is_final_op = dag.bfs_successors(*node).all(|(_, child_successors)| { + !child_successors.iter().any(|suc| match dag.dag()[*suc] { + NodeType::Operation(ref suc_inst) => { + !FINAL_OP_NAMES.contains(&suc_inst.op.name()) + } + _ => false, + }) + }); + is_final_op + }) + .collect(); + if final_ops.is_empty() { + return Ok(()); + } + let ordered_node_indices: Vec = dag + .topological_op_nodes()? + .filter(|node| final_ops.contains(node)) + .collect(); + let final_packed_ops: Vec = ordered_node_indices + .into_iter() + .map(|node| { + let NodeType::Operation(ref inst) = dag.dag()[node] else { + unreachable!() + }; + let res = inst.clone(); + dag.remove_op_node(node); + res + }) + .collect(); + let new_barrier = BARRIER + .get_bound(py) + .call1((dag.num_qubits(), label.as_deref()))?; + + let new_barrier_py_inst = PyInstruction { + qubits: dag.num_qubits() as u32, + clbits: 0, + params: 0, + op_name: "barrier".to_string(), + control_flow: false, + #[cfg(feature = "cache_pygates")] + instruction: new_barrier.clone().unbind(), + #[cfg(not(feature = "cache_pygates"))] + instruction: new_barrier.unbind(), + }; + let qargs: Vec = (0..dag.num_qubits() as u32).map(Qubit).collect(); + #[cfg(feature = "cache_pygates")] + { + dag.apply_operation_back( + py, + PackedOperation::from_instruction(Box::new(new_barrier_py_inst)), + qargs.as_slice(), + &[], + None, + ExtraInstructionAttributes::new(label, None, None, None), + Some(new_barrier.unbind()), + )?; + } + #[cfg(not(feature = "cache_pygates"))] + { + dag.apply_operation_back( + py, + PackedOperation::from_instruction(Box::new(new_barrier_py_inst)), + qargs.as_slice(), + &[], + None, + ExtraInstructionAttributes::new(label, None, None, None), + )?; + } + for inst in final_packed_ops { + dag.push_back(py, inst)?; + } + Ok(()) +} + +pub fn barrier_before_final_measurements_mod(m: &Bound) -> PyResult<()> { + m.add_wrapped(wrap_pyfunction!(barrier_before_final_measurements))?; + Ok(()) +} diff --git a/crates/accelerate/src/lib.rs b/crates/accelerate/src/lib.rs index f17fea5bccfe..a0a2bfcdcaad 100644 --- a/crates/accelerate/src/lib.rs +++ b/crates/accelerate/src/lib.rs @@ -14,6 +14,7 @@ use std::env; use pyo3::import_exception; +pub mod barrier_before_final_measurement; pub mod check_map; pub mod circuit_library; pub mod commutation_analysis; diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index 54f270b30558..31cc67d25e01 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -5074,7 +5074,7 @@ impl DAGCircuit { /// This is mostly used to apply operations from one DAG to /// another that was created from the first via /// [DAGCircuit::copy_empty_like]. - fn push_back(&mut self, py: Python, instr: PackedInstruction) -> PyResult { + pub fn push_back(&mut self, py: Python, instr: PackedInstruction) -> PyResult { let op_name = instr.op.name(); let (all_cbits, vars): (Vec, Option>) = { if self.may_have_additional_wires(py, &instr) { diff --git a/crates/pyext/src/lib.rs b/crates/pyext/src/lib.rs index 8b54c535db2f..dfda80623d6f 100644 --- a/crates/pyext/src/lib.rs +++ b/crates/pyext/src/lib.rs @@ -28,6 +28,7 @@ where #[rustfmt::skip] #[pymodule] fn _accelerate(m: &Bound) -> PyResult<()> { + add_submodule(m, ::qiskit_accelerate::barrier_before_final_measurement::barrier_before_final_measurements_mod, "barrier_before_final_measurement")?; add_submodule(m, ::qiskit_accelerate::check_map::check_map_mod, "check_map")?; add_submodule(m, ::qiskit_accelerate::circuit_library::circuit_library, "circuit_library")?; add_submodule(m, ::qiskit_accelerate::commutation_analysis::commutation_analysis, "commutation_analysis")?; diff --git a/qiskit/__init__.py b/qiskit/__init__.py index 0e9694217cd8..e96045e6801f 100644 --- a/qiskit/__init__.py +++ b/qiskit/__init__.py @@ -83,6 +83,9 @@ sys.modules["qiskit._accelerate.synthesis.permutation"] = _accelerate.synthesis.permutation sys.modules["qiskit._accelerate.synthesis.linear"] = _accelerate.synthesis.linear sys.modules["qiskit._accelerate.synthesis.clifford"] = _accelerate.synthesis.clifford +sys.modules["qiskit._accelerate.barrier_before_final_measurement"] = ( + _accelerate.barrier_before_final_measurement +) sys.modules["qiskit._accelerate.commutation_checker"] = _accelerate.commutation_checker sys.modules["qiskit._accelerate.commutation_analysis"] = _accelerate.commutation_analysis sys.modules["qiskit._accelerate.commutation_cancellation"] = _accelerate.commutation_cancellation diff --git a/qiskit/transpiler/passes/layout/disjoint_utils.py b/qiskit/transpiler/passes/layout/disjoint_utils.py index 42f665c473a1..3ecdc8137cea 100644 --- a/qiskit/transpiler/passes/layout/disjoint_utils.py +++ b/qiskit/transpiler/passes/layout/disjoint_utils.py @@ -107,7 +107,7 @@ def split_barriers(dag: DAGCircuit): num_qubits = len(node.qargs) if num_qubits == 1: continue - if node.op.label: + if node.label: barrier_uuid = f"{node.op.label}_uuid={uuid.uuid4()}" else: barrier_uuid = f"_none_uuid={uuid.uuid4()}" @@ -125,28 +125,30 @@ def split_barriers(dag: DAGCircuit): def combine_barriers(dag: DAGCircuit, retain_uuid: bool = True): """Mutate input dag to combine barriers with UUID labels into a single barrier.""" qubit_indices = {bit: index for index, bit in enumerate(dag.qubits)} - uuid_map: dict[uuid.UUID, DAGOpNode] = {} + uuid_map: dict[str, DAGOpNode] = {} for node in dag.op_nodes(Barrier): - if node.op.label: - if "_uuid=" in node.op.label: - barrier_uuid = node.op.label + if node.label: + if "_uuid=" in node.label: + barrier_uuid = node.label else: continue if barrier_uuid in uuid_map: other_node = uuid_map[barrier_uuid] num_qubits = len(other_node.qargs) + len(node.qargs) - new_op = Barrier(num_qubits, label=barrier_uuid) + if not retain_uuid: + if isinstance(node.label, str) and node.label.startswith("_none_uuid="): + label = None + elif isinstance(node.label, str) and "_uuid=" in node.label: + label = "_uuid=".join(node.label.split("_uuid=")[:-1]) + else: + label = barrier_uuid + else: + label = barrier_uuid + new_op = Barrier(num_qubits, label=label) new_node = dag.replace_block_with_op([node, other_node], new_op, qubit_indices) uuid_map[barrier_uuid] = new_node else: uuid_map[barrier_uuid] = node - if not retain_uuid: - for node in dag.op_nodes(Barrier): - if isinstance(node.op.label, str) and node.op.label.startswith("_none_uuid="): - node.op.label = None - elif isinstance(node.op.label, str) and "_uuid=" in node.op.label: - original_label = "_uuid=".join(node.op.label.split("_uuid=")[:-1]) - node.op.label = original_label def require_layout_isolated_to_component( diff --git a/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py b/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py index 4633cc57af54..c9270193f84e 100644 --- a/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py +++ b/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py @@ -13,9 +13,8 @@ """Add a barrier before final measurements.""" -from qiskit.circuit.barrier import Barrier +from qiskit._accelerate.barrier_before_final_measurement import barrier_before_final_measurements from qiskit.transpiler.basepasses import TransformationPass -from qiskit.dagcircuit import DAGCircuit, DAGOpNode from .merge_adjacent_barriers import MergeAdjacentBarriers @@ -33,60 +32,7 @@ def __init__(self, label=None): def run(self, dag): """Run the BarrierBeforeFinalMeasurements pass on `dag`.""" - # Collect DAG nodes which are followed only by barriers or other measures. - final_op_types = ["measure", "barrier"] - final_ops = [] - for candidate_node in dag.named_nodes(*final_op_types): - is_final_op = True - - for _, child_successors in dag.bfs_successors(candidate_node): - - if any( - isinstance(suc, DAGOpNode) and suc.name not in final_op_types - for suc in child_successors - ): - is_final_op = False - break - - if is_final_op: - final_ops.append(candidate_node) - - if not final_ops: - return dag - - # Create a layer with the barrier and add both bits and registers from the original dag. - barrier_layer = DAGCircuit() - barrier_layer.add_qubits(dag.qubits) - for qreg in dag.qregs.values(): - barrier_layer.add_qreg(qreg) - barrier_layer.add_clbits(dag.clbits) - for creg in dag.cregs.values(): - barrier_layer.add_creg(creg) - - # Add a barrier across all qubits so swap mapper doesn't add a swap - # from an unmeasured qubit after a measure. - final_qubits = dag.qubits - - barrier_layer.apply_operation_back( - Barrier(len(final_qubits), label=self.label), final_qubits, (), check=False - ) - - # Preserve order of final ops collected earlier from the original DAG. - ordered_final_nodes = [ - node for node in dag.topological_op_nodes() if node in set(final_ops) - ] - - # Move final ops to the new layer and append the new layer to the DAG. - for final_node in ordered_final_nodes: - barrier_layer.apply_operation_back( - final_node.op, final_node.qargs, final_node.cargs, check=False - ) - - for final_op in final_ops: - dag.remove_op_node(final_op) - - dag.compose(barrier_layer) - + barrier_before_final_measurements(dag, self.label) if self.label is None: # Merge the new barrier into any other barriers adjacent_pass = MergeAdjacentBarriers()