Add CriticalHeuristic to SABRE routing heuristics

crates/accelerate/src/sabre/heuristic.rs

@@ -165,6 +165,45 @@ impl DecayHeuristic {
     }
 }
 
+/// Define the characteristics of the "critical" heuristic. This is used to prioritize the gates in the critical path.
+#[pyclass]
+#[pyo3(module = "qiskit._accelerate.sabre", frozen)]
+#[derive(Clone, Copy, PartialEq)]
+pub struct CriticalHeuristic {
+    /// The relative weighting of this heuristic to others.  Typically you should just set this to
+    /// 1.0 and define everything else in terms of this.
+    pub weight: f64,
+    /// Set the dynamic scaling of the weight based on the layer it is applying to.
+    pub scale: SetScaling,
+}
+
+#[pymethods]
+impl CriticalHeuristic {
+    #[new]
+    pub fn new(weight: f64, scale: SetScaling) -> Self {
+        Self { weight, scale }
+    }
+
+    pub fn __getnewargs__(&self, py: Python) -> Py<PyAny> {
+        (self.weight, self.scale).into_py(py)
+    }
+
+    pub fn __eq__(&self, py: Python, other: Py<PyAny>) -> bool {
+        if let Ok(other) = other.extract::<Self>(py) {
+            self == &other
+        } else {
+            false
+        }
+    }
+
+    pub fn __repr__(&self, py: Python) -> PyResult<Py<PyAny>> {
+        let fmt = "CriticalHeuristic(weight={!r}, scale={!r})";
+        Ok(PyString::new_bound(py, fmt)
+            .call_method1("format", (self.weight, self.scale))?
+            .into_py(py))
+    }
+}
+
 /// A complete description of the heuristic that Sabre will use.  See the individual elements for a
 /// greater description.
 #[pyclass]
@@ -174,6 +213,7 @@ pub struct Heuristic {
     pub basic: Option<BasicHeuristic>,
     pub lookahead: Option<LookaheadHeuristic>,
     pub decay: Option<DecayHeuristic>,
+    pub critical: Option<CriticalHeuristic>,
     pub best_epsilon: f64,
     pub attempt_limit: usize,
 }
@@ -194,18 +234,20 @@ impl Heuristic {
     ///     best_epsilon (float): the floating-point epsilon to use when comparing scores to find
     ///         the best value.
     #[new]
-    #[pyo3(signature = (basic=None, lookahead=None, decay=None, attempt_limit=1000, best_epsilon=1e-10))]
+    #[pyo3(signature = (basic=None, lookahead=None, decay=None, critical=None,  attempt_limit=1000, best_epsilon=1e-10))]
     pub fn new(
         basic: Option<BasicHeuristic>,
         lookahead: Option<LookaheadHeuristic>,
         decay: Option<DecayHeuristic>,
+        critical: Option<CriticalHeuristic>,
         attempt_limit: Option<usize>,
         best_epsilon: f64,
     ) -> Self {
         Self {
             basic,
             lookahead,
             decay,
+            critical,
             best_epsilon,
             attempt_limit: attempt_limit.unwrap_or(usize::MAX),
         }
@@ -216,6 +258,7 @@ impl Heuristic {
             self.basic,
             self.lookahead,
             self.decay,
+            self.critical,
             self.attempt_limit,
             self.best_epsilon,
         )
@@ -258,6 +301,14 @@ impl Heuristic {
         }
     }
 
+    /// Set the multiplier increment and reset interval of the critical heuristic.
+    pub fn with_critical(&self, weight: f64, scale: SetScaling) -> Self {
+        Self {
+            critical: Some(CriticalHeuristic { weight, scale }),
+            ..self.clone()
+        }
+    }
+
     pub fn __eq__(&self, py: Python, other: Py<PyAny>) -> bool {
         if let Ok(other) = other.extract::<Self>(py) {
             self == &other
@@ -275,6 +326,7 @@ impl Heuristic {
                     self.basic,
                     self.lookahead,
                     self.decay,
+                    self.critical,
                     self.attempt_limit,
                     self.best_epsilon,
                 ),

crates/accelerate/src/sabre/layer.rs

@@ -26,6 +26,8 @@ use crate::nlayout::PhysicalQubit;
 /// makes it more efficient to do everything in terms of physical qubits, so the conversion between
 /// physical and virtual qubits via the layout happens once per inserted swap and on layer
 /// extension, not for every swap trialled.
+
+#[derive(Clone)]
 pub struct FrontLayer {
     /// Map of the (index to the) node to the qubits it acts on.
     nodes: IndexMap<NodeIndex, [PhysicalQubit; 2], ::ahash::RandomState>,

crates/accelerate/src/sabre/route.rs

@@ -32,7 +32,9 @@ use smallvec::{smallvec, SmallVec};
 use crate::getenv_use_multiple_threads;
 use crate::nlayout::{NLayout, PhysicalQubit};
 
-use super::heuristic::{BasicHeuristic, DecayHeuristic, Heuristic, LookaheadHeuristic, SetScaling};
+use super::heuristic::{
+    BasicHeuristic, CriticalHeuristic, DecayHeuristic, Heuristic, LookaheadHeuristic, SetScaling,
+};
 use super::layer::{ExtendedSet, FrontLayer};
 use super::neighbor_table::NeighborTable;
 use super::sabre_dag::SabreDAG;
@@ -69,6 +71,9 @@ struct RoutingState<'a, 'b> {
     front_layer: FrontLayer,
     extended_set: ExtendedSet,
     decay: &'a mut [f64],
+    /// Map from the node id to the ranking of the node in terms of the number of descendants.
+    /// Ranking of 1 means the node has the most descendants.
+    descendants_rank: HashMap<NodeIndex, usize>,
     /// How many predecessors still need to be satisfied for each node index before it is at the
     /// front of the topological iteration through the nodes as they're routed.
     required_predecessors: &'a mut [u32],
@@ -281,6 +286,39 @@ impl<'a, 'b> RoutingState<'a, 'b> {
         }
     }
 
+    /// Populate the `descendants_rank` map for the critical heuristic.
+    ///
+    /// This function calculates the number of descendants for each node in the DAG,
+    /// then ranks the nodes by the number of descendants, assigning a rank where
+    /// 1 represents the node with the most descendants.
+    /// The results are stored in the `descendants_rank` map, where each node ID maps to its rank.
+    /// This map is used by the critical heuristic to prioritize routing decisions.
+    fn populate_descendants_rank_map(&mut self) {
+        let mut node_id_to_descendants: HashMap<NodeIndex, usize> = HashMap::new();
+
+        // First, populate the number of descendants for each node.
+        for node in self.dag.dag.node_indices() {
+            let mut stack = vec![node];
+            let mut count = 0;
+            while let Some(n) = stack.pop() {
+                count += 1;
+                for edge in self.dag.dag.edges(n) {
+                    stack.push(edge.target());
+                }
+            }
+            node_id_to_descendants.insert(node, count);
+        }
+
+        // Sort nodes by the number of descendants and assign rankings.
+        let mut desc_list: Vec<_> = node_id_to_descendants.iter().collect();
+        desc_list.sort_by(|a, b| b.1.cmp(a.1)); // Sort in descending order of descendants
+
+        // Populate the `descendants_rank` map with rankings.
+        for (rank, (node_id, _)) in desc_list.into_iter().enumerate() {
+            self.descendants_rank.insert(*node_id, rank + 1);
+        }
+    }
+
     /// Add swaps to the current set that greedily bring the nearest node together.  This is a
     /// "release valve" mechanism; it ignores all the Sabre heuristics and forces progress, so we
     /// can't get permanently stuck.
@@ -416,6 +454,60 @@ impl<'a, 'b> RoutingState<'a, 'b> {
             }
         }
 
+        if let Some(CriticalHeuristic { weight, scale }) = self.heuristic.critical {
+            let weight = match scale {
+                SetScaling::Constant => weight,
+                SetScaling::Size => {
+                    if self.descendants_rank.is_empty() {
+                        0.0
+                    } else {
+                        weight / (self.descendants_rank.len() as f64)
+                    }
+                }
+            };
+
+            for (swap, score) in self.swap_scores.iter_mut() {
+                // Check what gates can be routed after the swap
+                let mut trial_front_layer = self.front_layer.clone();
+                trial_front_layer.apply_swap(*swap);
+
+                let mut routable_nodes = Vec::<NodeIndex>::with_capacity(2);
+                if let Some(node) =
+                    trial_front_layer.qubits()[swap[0].index()].and_then(|(node, other)| {
+                        self.target
+                            .coupling
+                            .contains_edge(
+                                NodeIndex::new(swap[0].index()),
+                                NodeIndex::new(other.index()),
+                            )
+                            .then_some(node)
+                    })
+                {
+                    routable_nodes.push(node);
+                }
+                if let Some(node) =
+                    trial_front_layer.qubits()[swap[1].index()].and_then(|(node, other)| {
+                        self.target
+                            .coupling
+                            .contains_edge(
+                                NodeIndex::new(swap[1].index()),
+                                NodeIndex::new(other.index()),
+                            )
+                            .then_some(node)
+                    })
+                {
+                    routable_nodes.push(node);
+                }
+
+                // For each routable node, substract 10^{-rank} from the score to prioritize routing in tie cases
+                for node in routable_nodes {
+                    if let Some(rank) = self.descendants_rank.get(&node) {
+                        *score -= weight / 10.0_f64.powi(*rank as i32);
+                    }
+                }
+            }
+        }
+
         let mut min_score = f64::INFINITY;
         let epsilon = self.heuristic.best_epsilon;
         for &(swap, score) in self.swap_scores.iter() {
@@ -553,6 +645,7 @@ pub fn swap_map_trial(
         front_layer: FrontLayer::new(num_qubits),
         extended_set: ExtendedSet::new(num_qubits),
         decay: &mut vec![1.; num_qubits as usize],
+        descendants_rank: HashMap::new(),
         required_predecessors: &mut vec![0; dag.dag.node_count()],
         layout: initial_layout.clone(),
         swap_scores: Vec::with_capacity(target.coupling.edge_count()),
@@ -565,6 +658,11 @@ pub fn swap_map_trial(
             state.required_predecessors[edge.target().index()] += 1;
         }
     }
+
+    if heuristic.critical.is_some() {
+        state.populate_descendants_rank_map();
+    }
+
     state.route_reachable_nodes(&dag.first_layer);
     state.populate_extended_set();
 

qiskit/transpiler/passes/routing/sabre_swap.py

@@ -229,6 +229,12 @@ def run(self, dag):
                 .with_lookahead(0.5, 20, SetScaling.Size)
                 .with_decay(0.001, 5)
             )
+        elif self.heuristic == "critical":
+            heuristic = (
+                Heuristic(attempt_limit=10 * num_dag_qubits)
+                .with_basic(1.0, SetScaling.Size)
+                .with_critical(0.5, SetScaling.Size)
+            )
         else:
             raise TranspilerError(f"Heuristic {self.heuristic} not recognized.")
         disjoint_utils.require_layout_isolated_to_component(

releasenotes/notes/sabre-heuristic-critical-308cfc1371eeaefe.yaml

@@ -0,0 +1,24 @@
+---
+prelude: >
+    A new feature has been added to the `SabreSwap` pass in Qiskit that introduces 
+    a `CriticalHeuristic`. This heuristic allows users to prioritize routing 
+    decisions based on the criticality of operations, as determined by the number 
+    of descendant gates in the quantum circuit's Directed Acyclic Graph (DAG). 
+    This enhancement aims to minimize circuit depth by ensuring that more critical 
+    operations, which have the most dependent operations, are routed first.
+
+features_transpiler:
+  - |
+    The `CriticalHeuristic` is a new addition to the `SabreSwap` pass in Qiskit. 
+    This heuristic ranks quantum gates based on their criticality, which is determined 
+    by the number of descendant gates they have in the circuit's DAG. Gates with more 
+    descendants are considered more critical and are prioritized during the routing process.
+
+    This heuristic works similarly to the existing ones like "basic," "lookahead," or "decay," 
+    but with a focus on ensuring that the critical path of the quantum circuit is minimized. 
+    This can lead to more efficient routing by reducing the potential for bottlenecks in 
+    the circuit, potentially lowering overall circuit depth.
+
+    This new heuristic gives users additional control over the routing process, enabling 
+    them to optimize their circuits by focusing on the critical paths, which can be 
+    particularly beneficial for complex quantum circuits with many dependencies.

test/python/transpiler/test_sabre_swap.py

@@ -247,7 +247,7 @@ def test_do_not_reorder_measurements(self):
         self.assertNotEqual(last_h.qubits, second_measure.qubits)
 
     # The 'basic' method can't get stuck in the same way.
-    @ddt.data("lookahead", "decay")
+    @ddt.data("lookahead", "decay", "critical")
     def test_no_infinite_loop(self, method):
         """Test that the 'release value' mechanisms allow SabreSwap to make progress even on
         circuits that get stuck in a stable local minimum of the lookahead parameters."""
@@ -371,7 +371,7 @@ def test_conditional_measurement(self):
         result = SabreSwap(CouplingMap.from_line(3), seed=12345)(qc)
         self.assertEqual(result, expected)
 
-    @ddt.data("basic", "lookahead", "decay")
+    @ddt.data("basic", "lookahead", "decay", "critical")
     def test_deterministic(self, heuristic):
         """Test that the output of the SabreSwap pass is deterministic for a given random seed."""
         width = 40