Perform BarrierBeforeFinalMeasurements analysis in parallel

crates/accelerate/src/barrier_before_final_measurement.rs

@@ -10,8 +10,8 @@
 // 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 rayon::prelude::*;
 use rustworkx_core::petgraph::stable_graph::NodeIndex;
 
 use qiskit_circuit::circuit_instruction::ExtraInstructionAttributes;
@@ -21,7 +21,7 @@ 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"];
+const PARALLEL_THRESHOLD: usize = 150;
 
 #[pyfunction]
 #[pyo3(signature=(dag, label=None))]
@@ -30,39 +30,87 @@ pub fn barrier_before_final_measurements(
     dag: &mut DAGCircuit,
     label: Option<String>,
 ) -> PyResult<()> {
-    let is_exactly_final = |inst: &PackedInstruction| FINAL_OP_NAMES.contains(&inst.op.name());
-    let final_ops: HashSet<NodeIndex> = dag
-        .op_nodes(true)
-        .filter_map(|(node, inst)| {
-            if !is_exactly_final(inst) {
-                return None;
-            }
-            dag.bfs_successors(node)
-                .all(|(_, child_successors)| {
-                    child_successors.iter().all(|suc| match dag[*suc] {
-                        NodeType::Operation(ref suc_inst) => is_exactly_final(suc_inst),
+    let find_final_nodes = |[_in_index, out_index]: &[NodeIndex; 2]| -> Vec<NodeIndex> {
+        let mut next_nodes: Vec<NodeIndex> = dag
+            .quantum_predecessors(*out_index)
+            .filter(|index| {
+                let node = &dag[*index];
+                match node {
+                    NodeType::Operation(inst) => {
+                        if inst.op.name() == "measure" || inst.op.name() == "barrier" {
+                            dag.bfs_successors(*index).all(|(_, child_successors)| {
+                                child_successors.iter().all(|suc| match &dag[*suc] {
+                                    NodeType::Operation(suc_inst) => {
+                                        suc_inst.op.name() == "measure"
+                                            || suc_inst.op.name() == "barrier"
+                                    }
+                                    _ => true,
+                                })
+                            })
+                        } else {
+                            false
+                        }
+                    }
+                    _ => false,
+                }
+            })
+            .collect();
+        let mut nodes: Vec<NodeIndex> = Vec::new();
+        while let Some(node_index) = next_nodes.pop() {
+            if node_index != *out_index
+                && dag.bfs_successors(node_index).all(|(_, child_successors)| {
+                    child_successors.iter().all(|suc| match &dag[*suc] {
+                        NodeType::Operation(suc_inst) => {
+                            suc_inst.op.name() == "measure" || suc_inst.op.name() == "barrier"
+                        }
                         _ => true,
                     })
                 })
-                .then_some(node)
-        })
-        .collect();
+            {
+                nodes.push(node_index);
+            }
+            for pred in dag.quantum_predecessors(node_index) {
+                match &dag[pred] {
+                    NodeType::Operation(inst) => {
+                        if inst.op.name() == "measure" || inst.op.name() == "barrier" {
+                            next_nodes.push(pred)
+                        }
+                    }
+                    _ => continue,
+                }
+            }
+        }
+        nodes.reverse();
+        nodes
+    };
+
+    let final_ops: Vec<NodeIndex> =
+        if dag.num_qubits() >= PARALLEL_THRESHOLD && crate::getenv_use_multiple_threads() {
+            dag.qubit_io_map()
+                .par_iter()
+                .flat_map(find_final_nodes)
+                .collect()
+        } else {
+            dag.qubit_io_map()
+                .iter()
+                .flat_map(find_final_nodes)
+                .collect()
+        };
+
     if final_ops.is_empty() {
         return Ok(());
     }
-    let ordered_node_indices: Vec<NodeIndex> = dag
-        .topological_op_nodes()?
-        .filter(|node| final_ops.contains(node))
-        .collect();
-    let final_packed_ops: Vec<PackedInstruction> = ordered_node_indices
+    let final_packed_ops: Vec<PackedInstruction> = final_ops
         .into_iter()
-        .map(|node| {
-            let NodeType::Operation(ref inst) = dag[node] else {
-                unreachable!()
-            };
-            let res = inst.clone();
-            dag.remove_op_node(node);
-            res
+        .filter_map(|node| match dag.dag().node_weight(node) {
+            Some(weight) => {
+                let NodeType::Operation(_) = weight else {
+                    return None;
+                };
+                let res = dag.remove_op_node(node);
+                Some(res)
+            }
+            None => None,
         })
         .collect();
     let new_barrier = BARRIER

crates/circuit/src/dag_circuit.rs

@@ -59,8 +59,8 @@ use rustworkx_core::petgraph::visit::{
 };
 use rustworkx_core::petgraph::Incoming;
 use rustworkx_core::traversal::{
-    ancestors as core_ancestors, bfs_successors as core_bfs_successors,
-    descendants as core_descendants,
+    ancestors as core_ancestors, bfs_predecessors as core_bfs_predecessors,
+    bfs_successors as core_bfs_successors, descendants as core_descendants,
 };
 
 use std::cmp::Ordering;
@@ -4768,6 +4768,12 @@ def _format(operand):
 }
 
 impl DAGCircuit {
+    /// Returns an immutable view of the qubit io map
+    #[inline(always)]
+    pub fn qubit_io_map(&self) -> &[[NodeIndex; 2]] {
+        &self.qubit_io_map
+    }
+
     /// Returns an immutable view of the inner StableGraph managed by the circuit.
     #[inline(always)]
     pub fn dag(&self) -> &StableDiGraph<NodeType, Wire> {
@@ -5577,7 +5583,7 @@ impl DAGCircuit {
     /// Remove an operation node n.
     ///
     /// Add edges from predecessors to successors.
-    pub fn remove_op_node(&mut self, index: NodeIndex) {
+    pub fn remove_op_node(&mut self, index: NodeIndex) -> PackedInstruction {
         let mut edge_list: Vec<(NodeIndex, NodeIndex, Wire)> = Vec::new();
         for (source, in_weight) in self
             .dag
@@ -5602,6 +5608,7 @@ impl DAGCircuit {
             Some(NodeType::Operation(packed)) => {
                 let op_name = packed.op.name();
                 self.decrement_op(op_name);
+                packed
             }
             _ => panic!("Must be called with valid operation node!"),
         }
@@ -5626,6 +5633,15 @@ impl DAGCircuit {
         core_bfs_successors(&self.dag, node).filter(move |(_, others)| !others.is_empty())
     }
 
+    /// Returns an iterator of tuples of (DAGNode, [DAGNodes]) where the DAGNode is the current node
+    /// and [DAGNode] is its successors in  BFS order.
+    pub fn bfs_predecessors(
+        &self,
+        node: NodeIndex,
+    ) -> impl Iterator<Item = (NodeIndex, Vec<NodeIndex>)> + '_ {
+        core_bfs_predecessors(&self.dag, node).filter(move |(_, others)| !others.is_empty())
+    }
+
     fn pack_into(&mut self, py: Python, b: &Bound<PyAny>) -> Result<NodeType, PyErr> {
         Ok(if let Ok(in_node) = b.downcast::<DAGInNode>() {
             let in_node = in_node.borrow();