Convert Sabre's lookahead to layer-based tracking

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.
+    /// The relative weights of each sequential layer in this heuristic.  Typically each of these is
+    /// defined relative to the :class:`.BasicHeuristic`, which generally has its weight set to 1.0.
+    weights: Vec<f64>,
+    /// Dynamic scaling of the heuristic weight depending on the size of each 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

@@ -18,29 +18,32 @@ use qiskit_circuit::PhysicalQubit;
 
 use super::vec_map::VecMap;
 
-/// A container for the current non-routable parts of the front layer.  This only ever holds
-/// two-qubit gates; the only reason a 0q- or 1q operation can be unroutable is because it has an
-/// unsatisfied 2q predecessor, which disqualifies it from being in the front layer.
+/// A container for 2q gates in a layer that are yet to be routed.
+///
+/// The graph nodes in this structure always refer to `TwoQ` entries, since `Synchronize` nodes do
+/// not impact the scoring; they only affect when nodes are eligible to move forwards a layer, or to
+/// become routed.
 ///
 /// It would be more algorithmically natural for this struct to work in terms of virtual qubits,
 /// because then a swap insertion would not change the data contained.  However, for each swap we
 /// insert, we score tens or hundreds, yet the subsequent update only affects two qubits.  This
 /// 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 {
-            // 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.
+        Layer {
+            // This is the maximum capacity of the 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(
                 num_qubits as usize / 2,
                 ::ahash::RandomState::default(),
@@ -60,15 +63,15 @@ impl FrontLayer {
         &self.qubits
     }
 
-    /// Add a node into the front layer, with the two qubits it operates on.
+    /// Add a node into the layer, with the two qubits it operates on.
     pub fn insert(&mut self, index: NodeIndex, qubits: [PhysicalQubit; 2]) {
         let [a, b] = qubits;
         self.qubits[a] = Some((index, b));
         self.qubits[b] = Some((index, a));
         self.nodes.insert(index, qubits);
     }
 
-    /// Remove a node from the front layer.
+    /// Remove a node from the layer.
     pub fn remove(&mut self, index: &NodeIndex) {
         // The actual order in the indexmap doesn't matter as long as it's reproducible.
         // Swap-remove is more efficient than a full shift-remove.
@@ -80,6 +83,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 {
@@ -89,15 +100,15 @@ impl FrontLayer {
     /// Calculate the score _difference_ caused by this swap, compared to not making the swap.
     #[inline(always)]
     pub fn score(&self, swap: [PhysicalQubit; 2], dist: &ArrayView2<f64>) -> f64 {
-        // 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
-        // never consider the same gate in both `if let` branches, because if we did, the gate would
-        // already be routable.  It doesn't matter, though, because the two distances would be
-        // equal anyway, so not affect the score.
+        // At most there can be two affected gates in the 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.
         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] {
@@ -106,7 +117,7 @@ impl FrontLayer {
         total
     }
 
-    /// Calculate the total absolute of the current front layer on the given layer.
+    /// Calculate the total absolute score of the layer for the set layout.
     pub fn total_score(&self, dist: &ArrayView2<f64>) -> f64 {
         self.iter()
             .map(|(_, &[a, b])| dist[[a.index(), b.index()]])
@@ -158,96 +169,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());
-    }
-}