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

(cherry picked from commit 43d8372)

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,22 +197,20 @@ 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()
             .flat_map(move |(a_index, others)| {
                 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,