Convert Sabre's lookahead to layer-based tracking

The previous Sabre extended set was just the "next N" 2q gates
topologically on from the front layer, where Qiskit reliably used
`N = 20` ever since its introduction.  For small-width circuits (as were
common when the original Sabre paper was written, and when it was first
implemented in Qiskit), this could mean the extended set was reliably
several layers deep.  This could also be the case for star-like
circuits.  For the wider circuits in use now, at the 100q order of
magnitude, the 20-gate limit reliably means that denser circuits cannot
have their entire next layer considered by the lookahead set.

This commit modifies the lookahead heuristic to be based specifically on
layers.  This regularises much of the structure of the heuristic with
respect to circuit and target topology; we reliably "look ahead" by the
same "distance" as far as routing is concerned.  It comes with the
additional benefits:

- we can use the same `Layer` structure for both the front layer and the
  lookahead layers, which reduces the amount of scoring code

- the lookahead score of a swap can now affect at most two gates per
  layer, just like the front-layer scoring, and we can do this
  statically without loops

- we no longer risk "biasing" the lookahead heuristic in either case of
  long chains of dependent gates (e.g. a gate that has 10 predecessors
  weights the score the same as a gate with only 1) or wide circuits
  (some qubits have their next layer counted in the score, but others
  don't because the extended set reached capacity).

- applying a swap to the lookahead now has a time complexity that is
  constant per layer, regardless of the number of gates stored in it,
  whereas previously it was proportional to the number of gates stored
  (and the implementation in the parent of this commit is proportional
  to the number of qubits in the circuit).

This change alone is mostly a set up, which enables further
computational complexity improvements by modifying the lookahead layers
in place after a gate routes, rather than rebuilding them from scratch,
and subsequently only updating _swap scores_ based on routing changes,
rather than recalculating all from scratch.

Author: jakelishman

crates/cext/src/transpiler/passes/sabre_layout.rs

@@ -120,11 +120,13 @@ pub unsafe extern "C" fn qk_transpiler_pass_standalone_sabre_layout(
             1.0,
             heuristic::SetScaling::Constant,
         )),
-        Some(heuristic::LookaheadHeuristic::new(
-            0.5,
-            20,
-            heuristic::SetScaling::Size,
-        )),
+        Some(
+            heuristic::LookaheadHeuristic::new(
+                vec![0.5 / target.num_qubits.unwrap_or(20) as f64],
+                heuristic::SetScaling::Constant,
+            )
+            .expect("number of layers should be valid"),
+        ),
         Some(heuristic::DecayHeuristic::new(0.001, 5)),
         Some(10 * target.num_qubits.unwrap() as usize),
         1e-10,

crates/transpiler/src/passes/sabre/heuristic.rs

@@ -79,43 +79,60 @@ impl BasicHeuristic {
     }
 }
 
-/// 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.
+/// Define the characteristics of the lookahead heuristic.
 #[pyclass(module = "qiskit._accelerate.sabre", frozen, from_py_object)]
-#[derive(Clone, Copy, PartialEq, Debug)]
+#[derive(Clone, PartialEq, Debug)]
 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.
+    weights: Vec<f64>,
+    /// Dynamic scaling of the heuristic weight depending on the size of the layer.
     pub scale: SetScaling,
 }
-impl_intopyobject_for_copy_pyclass!(LookaheadHeuristic);
+impl LookaheadHeuristic {
+    /// Construct a new lookahead heuristic.
+    ///
+    /// Fails if `weights` has $2^{16}$ or more elements (these are layers - 65,536 is too many!).
+    pub fn new(weights: Vec<f64>, scale: SetScaling) -> Option<Self> {
+        let _: u16 = weights.len().try_into().ok()?;
+        Some(Self { weights, scale })
+    }
+
+    pub fn num_layers(&self) -> u16 {
+        self.weights
+            .len()
+            .try_into()
+            .expect("constructor enforces sufficiently few layers")
+    }
+
+    #[inline]
+    pub fn weights(&self) -> &[f64] {
+        &self.weights
+    }
+}
 #[pymethods]
 impl LookaheadHeuristic {
     #[new]
-    pub fn new(weight: f64, size: usize, scale: SetScaling) -> Self {
-        Self {
-            weight,
-            size,
-            scale,
-        }
+    pub fn py_new(weights: Vec<f64>, scale: SetScaling) -> PyResult<Self> {
+        Self::new(weights, scale)
+            .ok_or_else(|| PyValueError::new_err("must have fewer than 65,536 layers"))
     }
 
     pub fn __getnewargs__(&self, py: Python) -> PyResult<Py<PyAny>> {
-        (self.weight, self.size, self.scale).into_py_any(py)
+        (self.weights.as_slice().into_pyobject(py)?, self.scale).into_py_any(py)
     }
 
     pub fn __eq__(&self, py: Python, other: Py<PyAny>) -> bool {
         other.extract::<Self>(py).is_ok_and(|other| self == &other)
     }
 
     pub fn __repr__(&self, py: Python) -> PyResult<Py<PyAny>> {
-        let fmt = "LookaheadHeuristic(weight={!r}, size={!r}, scale={!r})";
+        let fmt = "LookaheadHeuristic(weights={!r}, scale={!r})";
         PyString::new(py, fmt)
-            .call_method1("format", (self.weight, self.size, self.scale))?
+            .call_method1(
+                "format",
+                (self.weights.as_slice().into_pyobject(py)?, self.scale),
+            )?
             .into_py_any(py)
     }
 }
@@ -159,6 +176,14 @@ impl DecayHeuristic {
 
 /// A complete description of the heuristic that Sabre will use.  See the individual elements for a
 /// greater description.
+///
+/// .. note::
+///
+///     This is an internal Qiskit object, not a formal part of the API.  You can use this for
+///     fine-grained control over the Sabre heuristic, including for research, but beware that the
+///     available options and configuration of it may change without warning in minor versions of
+///     Qiskit.  If you are doing research using this, be sure to pin the version of Qiskit in your
+///     requirements.
 #[pyclass(module = "qiskit._accelerate.sabre", frozen, eq, skip_from_py_object)]
 #[derive(Clone, Debug, PartialEq)]
 pub struct Heuristic {
@@ -205,7 +230,7 @@ impl Heuristic {
     pub fn __getnewargs__(&self, py: Python) -> PyResult<Py<PyAny>> {
         (
             self.basic,
-            self.lookahead,
+            self.lookahead.clone(),
             self.decay,
             self.attempt_limit,
             self.best_epsilon,
@@ -223,15 +248,11 @@ impl Heuristic {
         }
     }
 
-    /// 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 {
+    /// Set the layer weights of the ``lookahead`` heuristic.  The weight here should typically be
+    /// less than that of ``basic``.  The number of weights dictates the number of layers.
+    pub fn with_lookahead(&self, weights: Vec<f64>, scale: SetScaling) -> Self {
         Self {
-            lookahead: Some(LookaheadHeuristic {
-                weight,
-                size,
-                scale,
-            }),
+            lookahead: Some(LookaheadHeuristic { weights, scale }),
             ..self.clone()
         }
     }
@@ -256,7 +277,7 @@ impl Heuristic {
                 "format",
                 (
                     self.basic,
-                    self.lookahead,
+                    self.lookahead.clone(),
                     self.decay,
                     self.attempt_limit,
                     self.best_epsilon,

crates/transpiler/src/passes/sabre/layer.rs

@@ -28,17 +28,18 @@ use super::vec_map::VecMap;
 /// 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.
-pub struct FrontLayer {
+#[derive(Clone, Debug)]
+pub struct Layer {
     /// Map of the (index to the) node to the qubits it acts on.
     nodes: IndexMap<NodeIndex, [PhysicalQubit; 2], ::ahash::RandomState>,
     /// Map of each qubit to the node that acts on it and the other qubit that node acts on, if this
     /// qubit is active (otherwise `None`).
     qubits: VecMap<PhysicalQubit, Option<(NodeIndex, PhysicalQubit)>>,
 }
 
-impl FrontLayer {
+impl Layer {
     pub fn new(num_qubits: u32) -> Self {
-        FrontLayer {
+        Layer {
             // This is the maximum capacity of the front layer, since each qubit must be one of a
             // pair, and can only have one gate in the layer.
             nodes: IndexMap::with_capacity_and_hasher(
@@ -80,6 +81,14 @@ impl FrontLayer {
         self.qubits[b] = None;
     }
 
+    /// Remove all nodes from the layer.
+    pub fn clear(&mut self) {
+        for (_, [a, b]) in self.nodes.drain(..) {
+            self.qubits[a] = None;
+            self.qubits[b] = None;
+        }
+    }
+
     /// Query whether a qubit has an active node.
     #[inline]
     pub fn is_active(&self, qubit: PhysicalQubit) -> bool {
@@ -98,6 +107,9 @@ impl FrontLayer {
         let [a, b] = swap;
         let mut total = 0.0;
         if let Some((_, c)) = self.qubits[a] {
+            if c == b {
+                return 0.0;
+            }
             total += dist[[b.index(), c.index()]] - dist[[a.index(), c.index()]]
         }
         if let Some((_, c)) = self.qubits[b] {
@@ -158,96 +170,3 @@ impl FrontLayer {
         self.nodes.values().flatten()
     }
 }
-
-/// This structure is currently reconstructed after each gate is routed, so there's no need to
-/// worry about tracking gate indices or anything like that.  We track length manually just to
-/// avoid a summation.
-pub struct ExtendedSet {
-    qubits: Vec<Vec<PhysicalQubit>>,
-    len: usize,
-}
-
-impl ExtendedSet {
-    pub fn new(num_qubits: u32) -> Self {
-        ExtendedSet {
-            qubits: vec![Vec::new(); num_qubits as usize],
-            len: 0,
-        }
-    }
-
-    /// Add a node and its active qubits to the extended set.
-    pub fn push(&mut self, qubits: [PhysicalQubit; 2]) {
-        let [a, b] = qubits;
-        self.qubits[a.index()].push(b);
-        self.qubits[b.index()].push(a);
-        self.len += 1;
-    }
-
-    /// 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 {
-        let [a, b] = swap;
-        let mut total = 0.0;
-        for other in self.qubits[a.index()].iter() {
-            // If the other qubit is also active then the score won't have changed, but since the
-            // distance is absolute, we'd double count rather than ignore if we didn't skip it.
-            if *other == b {
-                continue;
-            }
-            total += dist[[b.index(), other.index()]] - dist[[a.index(), other.index()]];
-        }
-        for other in self.qubits[b.index()].iter() {
-            if *other == a {
-                continue;
-            }
-            total += dist[[a.index(), other.index()]] - dist[[b.index(), other.index()]];
-        }
-        total
-    }
-
-    /// Calculate the total absolute score of this set of nodes over the given layout.
-    pub fn total_score(&self, dist: &ArrayView2<f64>) -> f64 {
-        // 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>()
-            * 0.5
-    }
-
-    /// Clear all nodes from the extended set.
-    pub fn clear(&mut self) {
-        for others in self.qubits.iter_mut() {
-            others.clear()
-        }
-        self.len = 0;
-    }
-
-    /// Number of nodes in the set.
-    pub fn len(&self) -> usize {
-        self.len
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.len == 0
-    }
-
-    /// Apply a physical swap to the current layout data structure.
-    pub fn apply_swap(&mut self, swap: [PhysicalQubit; 2]) {
-        let [a, b] = swap;
-        for other in self.qubits[a.index()].iter_mut() {
-            if *other == b {
-                *other = a
-            }
-        }
-        for other in self.qubits[b.index()].iter_mut() {
-            if *other == a {
-                *other = b
-            }
-        }
-        self.qubits.swap(a.index(), b.index());
-    }
-}