Expose Sabre heuristic configuration to Python (#12171)

* Expose Sabre heuristic configuration to Python

This exposes the entirety of the configuration of the Sabre heuristic to
Python space, making it modifiable without recompilation.  This includes
some additional configuration options that were not previously easily
modifiable, even with recompilation:

- the base weight of the "basic" component can be adjusted
- the weight of the "basic" and "lookahead" components can be adjusted
  to _either_ use a constant weight (previously not a thing) or use a
  weight that scales with the size of the set (previously the only
  option).
- the "decay" component is now entirely separated from the "lookahead"
  component, so in theory you can now have a decay without a lookahead.

This introduces a tracking `Vec` that stores the scores of _all_ the
swaps encountered, rather than just dynamically keeping hold of the best
swaps.  This has a couple of benefits:

- with the new dynamic structure for heuristics, this is rather more
  efficient because each heuristic component can be calculated in
  separate loops over the swaps, and we don't have to branch within the
  innermost loop.
- it makes it possible in the future to try things like assigning
  probabilities to each swap and randomly choosing from _all_ of them,
  not just the best swaps.  This is something I've actively wanted to
  try for quite some time.

The default heuristics in the transpiler-pass creators for the `basic`,
`lookahead` and `decay` strings are set to represent the same heuristics
as before, and this commit is entirely RNG compatible with its
predecessor (_technically_ for huge problems there's a possiblity that
pulling out some divisions into multiplications by reciprocals will
affect the floating-point maths enough to modify the swap selection).

* Update for PyO3 0.21

* Increase documentation of heuristic components

crates/accelerate/src/sabre/heuristic.rs

@@ -0,0 +1,284 @@
+// 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 pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyString;
+use pyo3::Python;
+
+/// Affect the dynamic scaling of the weight of node-set-based heuristics (basic and lookahead).
+#[pyclass]
+#[pyo3(module = "qiskit._accelerate.sabre", frozen)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum SetScaling {
+    /// No dynamic scaling of the weight.
+    Constant,
+    /// Scale the weight by the current number of nodes in the set (e.g., if it contains 5 nodes,
+    /// the weight will be multiplied by ``0.2``).
+    Size,
+}
+#[pymethods]
+impl SetScaling {
+    pub fn __reduce__(&self, py: Python) -> PyResult<Py<PyAny>> {
+        let name = match self {
+            SetScaling::Constant => "Constant",
+            SetScaling::Size => "Size",
+        };
+        Ok((
+            py.import_bound("builtins")?.getattr("getattr")?,
+            (py.get_type_bound::<Self>(), name),
+        )
+            .into_py(py))
+    }
+}
+
+/// Define the characteristics of the basic heuristic.  This is a simple sum of the physical
+/// distances of every gate in the front layer.
+#[pyclass]
+#[pyo3(module = "qiskit._accelerate.sabre", frozen)]
+#[derive(Clone, Copy, PartialEq)]
+pub struct BasicHeuristic {
+    /// 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 BasicHeuristic {
+    #[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 = "BasicHeuristic(weight={!r}, scale={!r})";
+        Ok(PyString::new_bound(py, fmt)
+            .call_method1("format", (self.weight, self.scale))?
+            .into_py(py))
+    }
+}
+
+/// Define the characteristics of the lookahead heuristic.  This is a sum of the physical distances
+/// of every gate in the lookahead set, which is gates immediately after the front layer.
+#[pyclass]
+#[pyo3(module = "qiskit._accelerate.sabre", frozen)]
+#[derive(Clone, Copy, PartialEq)]
+pub struct LookaheadHeuristic {
+    /// The relative weight of this heuristic.  Typically this is defined relative to the
+    /// :class:`.BasicHeuristic`, which generally has its weight set to 1.0.
+    pub weight: f64,
+    /// Number of gates to consider in the heuristic.
+    pub size: usize,
+    /// Dynamic scaling of the heuristic weight depending on the lookahead set.
+    pub scale: SetScaling,
+}
+#[pymethods]
+impl LookaheadHeuristic {
+    #[new]
+    pub fn new(weight: f64, size: usize, scale: SetScaling) -> Self {
+        Self {
+            weight,
+            size,
+            scale,
+        }
+    }
+
+    pub fn __getnewargs__(&self, py: Python) -> Py<PyAny> {
+        (self.weight, self.size, 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 = "LookaheadHeuristic(weight={!r}, size={!r}, scale={!r})";
+        Ok(PyString::new_bound(py, fmt)
+            .call_method1("format", (self.weight, self.size, self.scale))?
+            .into_py(py))
+    }
+}
+
+/// Define the characteristics of the "decay" heuristic.  In this, each physical qubit has a
+/// multiplier associated with it, beginning at 1.0, and has :attr:`increment` added to it each time
+/// the qubit is involved in a swap.  The final heuristic is calculated by multiplying all other
+/// components by the maximum multiplier involved in a given swap.
+#[pyclass]
+#[pyo3(module = "qiskit._accelerate.sabre", frozen)]
+#[derive(Clone, Copy, PartialEq)]
+pub struct DecayHeuristic {
+    /// The amount to add onto the multiplier of a physical qubit when it is used.
+    pub increment: f64,
+    /// How frequently (in terms of swaps in the layer) to reset all qubit multipliers back to 1.0.
+    pub reset: usize,
+}
+#[pymethods]
+impl DecayHeuristic {
+    #[new]
+    pub fn new(increment: f64, reset: usize) -> Self {
+        Self { increment, reset }
+    }
+
+    pub fn __getnewargs__(&self, py: Python) -> Py<PyAny> {
+        (self.increment, self.reset).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 = "DecayHeuristic(increment={!r}, reset={!r})";
+        Ok(PyString::new_bound(py, fmt)
+            .call_method1("format", (self.increment, self.reset))?
+            .into_py(py))
+    }
+}
+
+/// A complete description of the heuristic that Sabre will use.  See the individual elements for a
+/// greater description.
+#[pyclass]
+#[pyo3(module = "qiskit._accelerate.sabre", frozen)]
+#[derive(Clone, PartialEq)]
+pub struct Heuristic {
+    pub basic: Option<BasicHeuristic>,
+    pub lookahead: Option<LookaheadHeuristic>,
+    pub decay: Option<DecayHeuristic>,
+    pub best_epsilon: f64,
+    pub attempt_limit: usize,
+}
+
+#[pymethods]
+impl Heuristic {
+    /// Construct a new Sabre heuristic.  This can either be made directly of the desired
+    /// components, or you can make an empty heuristic and use the ``with_*`` methods to add
+    /// components to it.
+    ///
+    /// Args:
+    ///     attempt_limit (int): the maximum number of swaps to attempt before using a fallback
+    ///         "escape" mechanism to forcibly route a gate.  Set this to ``None`` to entirely
+    ///         disable the mechanism, but beware that it's possible (on large coupling maps with a
+    ///         lookahead heuristic component) for Sabre to get stuck in an inescapable arbitrarily
+    ///         deep local minimum of the heuristic.  If this happens, and the escape mechanism is
+    ///         disabled entirely, Sabre will enter an infinite loop.
+    ///     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))]
+    pub fn new(
+        basic: Option<BasicHeuristic>,
+        lookahead: Option<LookaheadHeuristic>,
+        decay: Option<DecayHeuristic>,
+        attempt_limit: Option<usize>,
+        best_epsilon: f64,
+    ) -> Self {
+        Self {
+            basic,
+            lookahead,
+            decay,
+            best_epsilon,
+            attempt_limit: attempt_limit.unwrap_or(usize::MAX),
+        }
+    }
+
+    pub fn __getnewargs__(&self, py: Python) -> Py<PyAny> {
+        (
+            self.basic,
+            self.lookahead,
+            self.decay,
+            self.attempt_limit,
+            self.best_epsilon,
+        )
+            .into_py(py)
+    }
+
+    /// Set the weight of the ``basic`` heuristic (the sum of distances of gates in the front
+    /// layer).  This is often set to ``1.0``.  You almost certainly should enable this part of the
+    /// heuristic, or it's highly unlikely that Sabre will be able to make any progress.
+    pub fn with_basic(&self, weight: f64, scale: SetScaling) -> Self {
+        Self {
+            basic: Some(BasicHeuristic { weight, scale }),
+            ..self.clone()
+        }
+    }
+
+    /// Set the weight and extended-set size of the ``lookahead`` heuristic.  The weight here
+    /// should typically be less than that of ``basic``.
+    pub fn with_lookahead(&self, weight: f64, size: usize, scale: SetScaling) -> Self {
+        Self {
+            lookahead: Some(LookaheadHeuristic {
+                weight,
+                size,
+                scale,
+            }),
+            ..self.clone()
+        }
+    }
+
+    /// Set the multiplier increment and reset interval of the decay heuristic.  The reset interval
+    /// must be non-zero.
+    pub fn with_decay(&self, increment: f64, reset: usize) -> PyResult<Self> {
+        if reset == 0 {
+            Err(PyValueError::new_err("decay reset interval cannot be zero"))
+        } else {
+            Ok(Self {
+                decay: Some(DecayHeuristic { increment, reset }),
+                ..self.clone()
+            })
+        }
+    }
+
+    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 = "Heuristic(basic={!r}, lookahead={!r}, decay={!r}, attempt_limit={!r}, best_epsilon={!r})";
+        Ok(PyString::new_bound(py, fmt)
+            .call_method1(
+                "format",
+                (
+                    self.basic,
+                    self.lookahead,
+                    self.decay,
+                    self.attempt_limit,
+                    self.best_epsilon,
+                ),
+            )?
+            .into_py(py))
+    }
+}

crates/accelerate/src/sabre/layer.rs

@@ -47,6 +47,11 @@ impl FrontLayer {
         }
     }
 
+    /// Number of gates currently stored in the layer.
+    pub fn len(&self) -> usize {
+        self.nodes.len()
+    }
+
     /// View onto the mapping between qubits and their `(node, other_qubit)` pair.  Index `i`
     /// corresponds to physical qubit `i`.
     pub fn qubits(&self) -> &[Option<(NodeIndex, PhysicalQubit)>] {
@@ -77,11 +82,8 @@ impl FrontLayer {
     }
 
     /// Calculate the score _difference_ caused by this swap, compared to not making the swap.
-    #[inline]
+    #[inline(always)]
     pub fn score(&self, swap: [PhysicalQubit; 2], dist: &ArrayView2<f64>) -> f64 {
-        if self.is_empty() {
-            return 0.0;
-        }
         // At most there can be two affected gates in the front layer (one on each qubit in the
         // swap), since any gate whose closest path passes through the swapped qubit link has its
         // "virtual-qubit path" order changed, but not the total weight.  In theory, we should
@@ -96,18 +98,14 @@ impl FrontLayer {
         if let Some((_, c)) = self.qubits[b.index()] {
             total += dist[[a.index(), c.index()]] - dist[[b.index(), c.index()]]
         }
-        total / self.nodes.len() as f64
+        total
     }
 
     /// Calculate the total absolute of the current front layer on the given layer.
     pub fn total_score(&self, dist: &ArrayView2<f64>) -> f64 {
-        if self.is_empty() {
-            return 0.0;
-        }
         self.iter()
             .map(|(_, &[a, b])| dist[[a.index(), b.index()]])
             .sum::<f64>()
-            / self.nodes.len() as f64
     }
 
     /// Apply a physical swap to the current layout data structure.
@@ -181,10 +179,8 @@ impl ExtendedSet {
     }
 
     /// Calculate the score of applying the given swap, relative to not applying it.
+    #[inline(always)]
     pub fn score(&self, swap: [PhysicalQubit; 2], dist: &ArrayView2<f64>) -> f64 {
-        if self.is_empty() {
-            return 0.0;
-        }
         let [a, b] = swap;
         let mut total = 0.0;
         for other in self.qubits[a.index()].iter() {
@@ -201,14 +197,12 @@ impl ExtendedSet {
             }
             total += dist[[a.index(), other.index()]] - dist[[b.index(), other.index()]];
         }
-        total / self.len as f64
+        total
     }
 
     /// Calculate the total absolute score of this set of nodes over the given layout.
     pub fn total_score(&self, dist: &ArrayView2<f64>) -> f64 {
-        if self.is_empty() {
-            return 0.0;
-        }
+        // Factor of two is to remove double-counting of each gate.
         self.qubits
             .iter()
             .enumerate()
@@ -216,7 +210,7 @@ impl ExtendedSet {
                 others.iter().map(move |b| dist[[a_index, b.index()]])
             })
             .sum::<f64>()
-            / (2.0 * self.len as f64) // Factor of two is to remove double-counting of each gate.
+            * 0.5
     }
 
     /// Clear all nodes from the extended set.

crates/accelerate/src/sabre/layout.rs

@@ -24,11 +24,12 @@ use rayon::prelude::*;
 use crate::getenv_use_multiple_threads;
 use crate::nlayout::{NLayout, PhysicalQubit};
 
+use super::heuristic::Heuristic;
 use super::neighbor_table::NeighborTable;
 use super::route::{swap_map, swap_map_trial, RoutingTargetView};
 use super::sabre_dag::SabreDAG;
 use super::swap_map::SwapMap;
-use super::{Heuristic, NodeBlockResults, SabreResult};
+use super::{NodeBlockResults, SabreResult};
 
 use crate::dense_layout::best_subset_inner;
 
@@ -39,7 +40,7 @@ pub fn sabre_layout_and_routing(
     dag: &SabreDAG,
     neighbor_table: &NeighborTable,
     distance_matrix: PyReadonlyArray2<f64>,
-    heuristic: Heuristic,
+    heuristic: &Heuristic,
     max_iterations: usize,
     num_swap_trials: usize,
     num_random_trials: usize,
@@ -129,7 +130,7 @@ pub fn sabre_layout_and_routing(
 fn layout_trial(
     target: &RoutingTargetView,
     dag: &SabreDAG,
-    heuristic: Heuristic,
+    heuristic: &Heuristic,
     seed: u64,
     max_iterations: usize,
     num_swap_trials: usize,

crates/accelerate/src/sabre/mod.rs

@@ -10,6 +10,7 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+mod heuristic;
 mod layer;
 mod layout;
 mod neighbor_table;
@@ -29,14 +30,6 @@ use neighbor_table::NeighborTable;
 use sabre_dag::SabreDAG;
 use swap_map::SwapMap;
 
-#[pyclass]
-#[derive(Clone, Copy)]
-pub enum Heuristic {
-    Basic,
-    Lookahead,
-    Decay,
-}
-
 /// A container for Sabre mapping results.
 #[pyclass(module = "qiskit._accelerate.sabre")]
 #[derive(Clone, Debug)]
@@ -117,7 +110,11 @@ impl BlockResult {
 pub fn sabre(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(route::sabre_routing))?;
     m.add_wrapped(wrap_pyfunction!(layout::sabre_layout_and_routing))?;
-    m.add_class::<Heuristic>()?;
+    m.add_class::<heuristic::SetScaling>()?;
+    m.add_class::<heuristic::Heuristic>()?;
+    m.add_class::<heuristic::BasicHeuristic>()?;
+    m.add_class::<heuristic::LookaheadHeuristic>()?;
+    m.add_class::<heuristic::DecayHeuristic>()?;
     m.add_class::<NeighborTable>()?;
     m.add_class::<SabreDAG>()?;
     m.add_class::<SwapMap>()?;

crates/accelerate/src/sabre/route.rs

@@ -31,22 +31,13 @@ use rustworkx_core::token_swapper::token_swapper;
 use crate::getenv_use_multiple_threads;
 use crate::nlayout::{NLayout, PhysicalQubit};
 
+use super::heuristic::{BasicHeuristic, DecayHeuristic, Heuristic, LookaheadHeuristic, SetScaling};
 use super::layer::{ExtendedSet, FrontLayer};
 use super::neighbor_table::NeighborTable;
 use super::sabre_dag::SabreDAG;
 use super::swap_map::SwapMap;
-use super::{BlockResult, Heuristic, NodeBlockResults, SabreResult};
+use super::{BlockResult, NodeBlockResults, SabreResult};
 
-/// Epsilon used in minimum-score calculations.
-const BEST_EPSILON: f64 = 1e-10;
-/// Size of lookahead window.
-const EXTENDED_SET_SIZE: usize = 20;
-/// Decay coefficient for penalizing serial swaps.
-const DECAY_RATE: f64 = 0.001;
-/// How often to reset all decay rates to 1.
-const DECAY_RESET_INTERVAL: u8 = 5;
-/// Weight of lookahead window compared to front_layer.
-const EXTENDED_SET_WEIGHT: f64 = 0.5;
 /// Number of trials for control flow block swap epilogues.
 const SWAP_EPILOGUE_TRIALS: usize = 4;
 
@@ -67,7 +58,7 @@ pub struct RoutingTargetView<'a> {
 struct RoutingState<'a, 'b> {
     target: &'a RoutingTargetView<'b>,
     dag: &'a SabreDAG,
-    heuristic: Heuristic,
+    heuristic: &'a Heuristic,
     /// Mapping of instructions (node indices) to swaps that precede them.
     out_map: HashMap<usize, Vec<[PhysicalQubit; 2]>>,
     /// Order of the instructions (node indices) in the problem DAG in the output.
@@ -76,16 +67,17 @@ struct RoutingState<'a, 'b> {
     node_block_results: HashMap<usize, Vec<BlockResult>>,
     front_layer: FrontLayer,
     extended_set: ExtendedSet,
+    decay: &'a mut [f64],
     /// 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],
     layout: NLayout,
-    /// Tracking for the 'decay' heuristic on each qubit.
-    qubits_decay: &'a mut [f64],
-    /// Reusable allocated storage space for choosing the best swap.  This is owned outside of the
-    /// `choose_best_swap` function so that we don't need to reallocate and then re-grow the
-    /// collection on every entry.
-    swap_scratch: Vec<[PhysicalQubit; 2]>,
+    /// Reusable allocated storage space for accumulating and scoring swaps.  This is owned as part
+    /// of the general state to avoid reallocation costs.
+    swap_scores: Vec<([PhysicalQubit; 2], f64)>,
+    /// Reusable allocated storage space for tracking the current best swaps.  This is owned as
+    /// part of the general state to avoid reallocation costs.
+    best_swaps: Vec<[PhysicalQubit; 2]>,
     rng: Pcg64Mcg,
     seed: u64,
 }
@@ -241,13 +233,19 @@ impl<'a, 'b> RoutingState<'a, 'b> {
     /// layer (and themselves).  This uses `required_predecessors` as scratch space for efficiency,
     /// but returns it to the same state as the input on return.
     fn populate_extended_set(&mut self) {
+        let extended_set_size =
+            if let Some(LookaheadHeuristic { size, .. }) = self.heuristic.lookahead {
+                size
+            } else {
+                return;
+            };
         let mut to_visit = self.front_layer.iter_nodes().copied().collect::<Vec<_>>();
         let mut decremented: IndexMap<usize, u32, ahash::RandomState> =
             IndexMap::with_hasher(ahash::RandomState::default());
         let mut i = 0;
         let mut visit_now: Vec<NodeIndex> = Vec::new();
         let dag = &self.dag;
-        while i < to_visit.len() && self.extended_set.len() < EXTENDED_SET_SIZE {
+        while i < to_visit.len() && self.extended_set.len() < extended_set_size {
             // Visit runs of non-2Q gates fully before moving on to children of 2Q gates. This way,
             // traversal order is a BFS of 2Q gates rather than of all gates.
             visit_now.push(to_visit[i]);
@@ -335,61 +333,76 @@ impl<'a, 'b> RoutingState<'a, 'b> {
 
     /// Return the swap of two virtual qubits that produces the best score of all possible swaps.
     fn choose_best_swap(&mut self) -> [PhysicalQubit; 2] {
-        self.swap_scratch.clear();
-        let mut min_score = f64::MAX;
-        // The decay heuristic is the only one that actually needs the absolute score.
-        let dist = &self.target.distance;
-        let absolute_score = match self.heuristic {
-            Heuristic::Decay => {
-                self.front_layer.total_score(dist)
-                    + EXTENDED_SET_WEIGHT * self.extended_set.total_score(dist)
-            }
-            _ => 0.0,
-        };
-        for swap in obtain_swaps(&self.front_layer, self.target.neighbors) {
-            let score = match self.heuristic {
-                Heuristic::Basic => self.front_layer.score(swap, dist),
-                Heuristic::Lookahead => {
-                    self.front_layer.score(swap, dist)
-                        + EXTENDED_SET_WEIGHT * self.extended_set.score(swap, dist)
+        // Obtain all candidate swaps from the front layer.  A coupling-map edge is a candidate
+        // swap if it involves at least one active qubit (i.e. it must affect the "basic"
+        // heuristic), and if it involves two active qubits, we choose the `swap[0] < swap[1]` form
+        // to make a canonical choice.
+        self.swap_scores.clear();
+        for &phys in self.front_layer.iter_active() {
+            for &neighbor in self.target.neighbors[phys].iter() {
+                if neighbor > phys || !self.front_layer.is_active(neighbor) {
+                    self.swap_scores.push(([phys, neighbor], 0.0));
                 }
-                Heuristic::Decay => {
-                    self.qubits_decay[swap[0].index()].max(self.qubits_decay[swap[1].index()])
-                        * (absolute_score
-                            + self.front_layer.score(swap, dist)
-                            + EXTENDED_SET_WEIGHT * self.extended_set.score(swap, dist))
-                }
-            };
-            if score < min_score - BEST_EPSILON {
-                min_score = score;
-                self.swap_scratch.clear();
-                self.swap_scratch.push(swap);
-            } else if (score - min_score).abs() < BEST_EPSILON {
-                self.swap_scratch.push(swap);
             }
         }
-        *self.swap_scratch.choose(&mut self.rng).unwrap()
-    }
-}
 
-/// Return a set of candidate swaps that affect qubits in front_layer.
-///
-/// For each virtual qubit in `front_layer`, find its current location on hardware and the physical
-/// qubits in that neighborhood. Every swap on virtual qubits that corresponds to one of those
-/// physical couplings is a candidate swap.
-fn obtain_swaps<'a>(
-    front_layer: &'a FrontLayer,
-    neighbors: &'a NeighborTable,
-) -> impl Iterator<Item = [PhysicalQubit; 2]> + 'a {
-    front_layer.iter_active().flat_map(move |&p| {
-        neighbors[p].iter().filter_map(move |&neighbor| {
-            if neighbor > p || !front_layer.is_active(neighbor) {
-                Some([p, neighbor])
-            } else {
-                None
+        let dist = &self.target.distance;
+        let mut absolute_score = 0.0;
+
+        if let Some(BasicHeuristic { weight, scale }) = self.heuristic.basic {
+            let weight = match scale {
+                SetScaling::Constant => weight,
+                SetScaling::Size => {
+                    if self.front_layer.is_empty() {
+                        0.0
+                    } else {
+                        weight / (self.front_layer.len() as f64)
+                    }
+                }
+            };
+            absolute_score += weight * self.front_layer.total_score(dist);
+            for (swap, score) in self.swap_scores.iter_mut() {
+                *score += weight * self.front_layer.score(*swap, dist);
             }
-        })
-    })
+        }
+
+        if let Some(LookaheadHeuristic { weight, scale, .. }) = self.heuristic.lookahead {
+            let weight = match scale {
+                SetScaling::Constant => weight,
+                SetScaling::Size => {
+                    if self.extended_set.is_empty() {
+                        0.0
+                    } else {
+                        weight / (self.extended_set.len() as f64)
+                    }
+                }
+            };
+            absolute_score += weight * self.extended_set.total_score(dist);
+            for (swap, score) in self.swap_scores.iter_mut() {
+                *score += weight * self.extended_set.score(*swap, dist);
+            }
+        }
+
+        if let Some(DecayHeuristic { .. }) = self.heuristic.decay {
+            for (swap, score) in self.swap_scores.iter_mut() {
+                *score = (absolute_score + *score)
+                    * self.decay[swap[0].index()].max(self.decay[swap[1].index()]);
+            }
+        }
+
+        let mut min_score = f64::INFINITY;
+        let epsilon = self.heuristic.best_epsilon;
+        for &(swap, score) in self.swap_scores.iter() {
+            if score - min_score < -epsilon {
+                min_score = score;
+                self.best_swaps.clear();
+                self.best_swaps.push(swap);
+            } else if (score - min_score).abs() <= epsilon {
+                self.best_swaps.push(swap);
+            }
+        }
+        *self.best_swaps.choose(&mut self.rng).unwrap()
+    }
 }
 
 /// Run sabre swap on a circuit
@@ -408,7 +421,7 @@ pub fn sabre_routing(
     dag: &SabreDAG,
     neighbor_table: &NeighborTable,
     distance_matrix: PyReadonlyArray2<f64>,
-    heuristic: Heuristic,
+    heuristic: &Heuristic,
     initial_layout: &NLayout,
     num_trials: usize,
     seed: Option<u64>,
@@ -449,7 +462,7 @@ pub fn sabre_routing(
 pub fn swap_map(
     target: &RoutingTargetView,
     dag: &SabreDAG,
-    heuristic: Heuristic,
+    heuristic: &Heuristic,
     initial_layout: &NLayout,
     seed: Option<u64>,
     num_trials: usize,
@@ -498,7 +511,7 @@ pub fn swap_map(
 pub fn swap_map_trial(
     target: &RoutingTargetView,
     dag: &SabreDAG,
-    heuristic: Heuristic,
+    heuristic: &Heuristic,
     initial_layout: &NLayout,
     seed: u64,
 ) -> (SabreResult, NLayout) {
@@ -512,10 +525,11 @@ pub fn swap_map_trial(
         node_block_results: HashMap::with_capacity(dag.node_blocks.len()),
         front_layer: FrontLayer::new(num_qubits),
         extended_set: ExtendedSet::new(num_qubits),
+        decay: &mut vec![1.; num_qubits as usize],
         required_predecessors: &mut vec![0; dag.dag.node_count()],
         layout: initial_layout.clone(),
-        qubits_decay: &mut vec![1.; num_qubits as usize],
-        swap_scratch: Vec::new(),
+        swap_scores: Vec::with_capacity(target.coupling.edge_count()),
+        best_swaps: Vec::new(),
         rng: Pcg64Mcg::seed_from_u64(seed),
         seed,
     };
@@ -529,15 +543,14 @@ pub fn swap_map_trial(
 
     // Main logic loop; the front layer only becomes empty when all nodes have been routed.  At
     // each iteration of this loop, we route either one or two gates.
-    let max_iterations_without_progress = 10 * num_qubits as usize;
-    let mut num_search_steps: u8 = 0;
     let mut routable_nodes = Vec::<NodeIndex>::with_capacity(2);
+    let mut num_search_steps = 0;
 
     while !state.front_layer.is_empty() {
         let mut current_swaps: Vec<[PhysicalQubit; 2]> = Vec::new();
         // Swap-mapping loop.  This is the main part of the algorithm, which we repeat until we
         // either successfully route a node, or exceed the maximum number of attempts.
-        while routable_nodes.is_empty() && current_swaps.len() <= max_iterations_without_progress {
+        while routable_nodes.is_empty() && current_swaps.len() <= state.heuristic.attempt_limit {
             let best_swap = state.choose_best_swap();
             state.apply_swap(best_swap);
             current_swaps.push(best_swap);
@@ -547,13 +560,15 @@ pub fn swap_map_trial(
             if let Some(node) = state.routable_node_on_qubit(best_swap[0]) {
                 routable_nodes.push(node);
             }
-            num_search_steps += 1;
-            if num_search_steps >= DECAY_RESET_INTERVAL {
-                state.qubits_decay.fill(1.);
-                num_search_steps = 0;
-            } else {
-                state.qubits_decay[best_swap[0].index()] += DECAY_RATE;
-                state.qubits_decay[best_swap[1].index()] += DECAY_RATE;
+            if let Some(DecayHeuristic { increment, reset }) = state.heuristic.decay {
+                num_search_steps += 1;
+                if num_search_steps >= reset {
+                    state.decay.fill(1.);
+                    num_search_steps = 0;
+                } else {
+                    state.decay[best_swap[0].index()] += increment;
+                    state.decay[best_swap[1].index()] += increment;
+                }
             }
         }
         if routable_nodes.is_empty() {
@@ -568,7 +583,9 @@ pub fn swap_map_trial(
             routable_nodes.push(force_routed);
         }
         state.update_route(&routable_nodes, current_swaps);
-        state.qubits_decay.fill(1.);
+        if state.heuristic.decay.is_some() {
+            state.decay.fill(1.);
+        }
         routable_nodes.clear();
     }
     (

qiskit/transpiler/passes/layout/sabre_layout.py

@@ -35,11 +35,7 @@ from qiskit.transpiler.layout import Layout
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.exceptions import TranspilerError
 from qiskit._accelerate.nlayout import NLayout
-from qiskit._accelerate.sabre import (
-    sabre_layout_and_routing,
-    Heuristic,
-    NeighborTable,
-)
+from qiskit._accelerate.sabre import sabre_layout_and_routing, Heuristic, NeighborTable, SetScaling
 from qiskit.transpiler.passes.routing.sabre_swap import _build_sabre_dag, _apply_sabre_result
 from qiskit.transpiler.target import Target
 from qiskit.transpiler.coupling import CouplingMap
@@ -393,12 +389,18 @@ class SabreLayout(TransformationPass):
             coupling_map.size(),
             original_qubit_indices,
         )
+        heuristic = (
+            Heuristic(attempt_limit=10 * coupling_map.size())
+            .with_basic(1.0, SetScaling.Size)
+            .with_lookahead(0.5, 20, SetScaling.Size)
+            .with_decay(0.001, 5)
+        )
         sabre_start = time.perf_counter()
         (initial_layout, final_permutation, sabre_result) = sabre_layout_and_routing(
             sabre_dag,
             neighbor_table,
             dist_matrix,
-            Heuristic.Decay,
+            heuristic,
             self.max_iterations,
             self.swap_trials,
             self.layout_trials,

qiskit/transpiler/passes/routing/sabre_swap.py

@@ -31,12 +31,7 @@ from qiskit.transpiler.passes.layout import disjoint_utils
 from qiskit.dagcircuit import DAGCircuit
 from qiskit.utils.parallel import CPU_COUNT
 
-from qiskit._accelerate.sabre import (
-    sabre_routing,
-    Heuristic,
-    NeighborTable,
-    SabreDAG,
-)
+from qiskit._accelerate.sabre import sabre_routing, Heuristic, SetScaling, NeighborTable, SabreDAG
 from qiskit._accelerate.nlayout import NLayout
 
 logger = logging.getLogger(__name__)
@@ -211,12 +206,25 @@ class SabreSwap(TransformationPass):
                 " This circuit cannot be routed to this device."
             )
 
-        if self.heuristic == "basic":
-            heuristic = Heuristic.Basic
+        if isinstance(self.heuristic, Heuristic):
+            heuristic = self.heuristic
+        elif self.heuristic == "basic":
+            heuristic = Heuristic(attempt_limit=10 * num_dag_qubits).with_basic(
+                1.0, SetScaling.Size
+            )
         elif self.heuristic == "lookahead":
-            heuristic = Heuristic.Lookahead
+            heuristic = (
+                Heuristic(attempt_limit=10 * num_dag_qubits)
+                .with_basic(1.0, SetScaling.Size)
+                .with_lookahead(0.5, 20, SetScaling.Size)
+            )
         elif self.heuristic == "decay":
-            heuristic = Heuristic.Decay
+            heuristic = (
+                Heuristic(attempt_limit=10 * num_dag_qubits)
+                .with_basic(1.0, SetScaling.Size)
+                .with_lookahead(0.5, 20, SetScaling.Size)
+                .with_decay(0.001, 5)
+            )
         else:
             raise TranspilerError(f"Heuristic {self.heuristic} not recognized.")
         disjoint_utils.require_layout_isolated_to_component(