From 23d996800f30a09aaa311e7057983a886ad4efaa Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Wed, 7 Feb 2024 12:41:40 -0500 Subject: [PATCH] Optimize BarrierBeforeFinalMeasurement pass The barrier before final measurement pass previously was working by iterating over the DAGCircuit to find all the barriers and measurements and then evaluating if those operations were at the end of the circuit, or adjacent to only barriers prior to the end of the circuit. However, this was fairly inefficient as it means the performance of the pass scales as a function of the number of gates in the circuit. This commit optimizes the performance as a function by looking at the predecessors of each qubit's output nodes to find final measurements instead of iterating over the entire circuit. --- .../barrier_before_final_measurements.py | 52 +++++++++++-------- .../test_barrier_before_final_measurements.py | 1 - 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py b/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py index 4633cc57af54..c1004bd86de2 100644 --- a/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py +++ b/qiskit/transpiler/passes/utils/barrier_before_final_measurements.py @@ -13,9 +13,12 @@ """Add a barrier before final measurements.""" +import rustworkx as rx + from qiskit.circuit.barrier import Barrier +from qiskit.circuit import Qubit from qiskit.transpiler.basepasses import TransformationPass -from qiskit.dagcircuit import DAGCircuit, DAGOpNode +from qiskit.dagcircuit import DAGCircuit, DAGOpNode, DAGOutNode from .merge_adjacent_barriers import MergeAdjacentBarriers @@ -34,24 +37,29 @@ 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: + final_op_types = {"measure", "barrier"} + final_ops_set = set() + for output_node in dag.output_map.values(): + qubit = output_node.wire + if not isinstance(qubit, Qubit): + continue + pred = next(dag.predecessors(output_node)) + + if ( + isinstance(pred, DAGOpNode) + and pred.name in final_op_types + and all(isinstance(x, DAGOutNode) for x in dag.successors(pred)) + ): + final_ops_set.add(pred) + for new_pred, parent_preds in rx.bfs_predecessors(dag._multi_graph, pred._node_id): + if any( + isinstance(p, DAGOpNode) and p.name not in final_op_types + for p in parent_preds + ): + break + final_ops_set.add(new_pred) + + if not final_ops_set: return dag # Create a layer with the barrier and add both bits and registers from the original dag. @@ -72,9 +80,7 @@ def run(self, dag): ) # 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) - ] + ordered_final_nodes = [node for node in dag.topological_op_nodes() if node in final_ops_set] # Move final ops to the new layer and append the new layer to the DAG. for final_node in ordered_final_nodes: @@ -82,7 +88,7 @@ def run(self, dag): final_node.op, final_node.qargs, final_node.cargs, check=False ) - for final_op in final_ops: + for final_op in ordered_final_nodes: dag.remove_op_node(final_op) dag.compose(barrier_layer) diff --git a/test/python/transpiler/test_barrier_before_final_measurements.py b/test/python/transpiler/test_barrier_before_final_measurements.py index eb4f509e5301..f664df3db92d 100644 --- a/test/python/transpiler/test_barrier_before_final_measurements.py +++ b/test/python/transpiler/test_barrier_before_final_measurements.py @@ -187,7 +187,6 @@ def test_preserve_measure_for_conditional(self): pass_ = BarrierBeforeFinalMeasurements() result = pass_.run(circuit_to_dag(circuit)) - self.assertEqual(result, circuit_to_dag(expected))