Add option to SabreLayout to specify starting layouts

crates/accelerate/src/sabre_layout.rs

@@ -11,6 +11,7 @@
 // that they have been altered from the originals.
 #![allow(clippy::too_many_arguments)]
 
+use hashbrown::HashSet;
 use ndarray::prelude::*;
 use numpy::IntoPyArray;
 use numpy::PyReadonlyArray2;
@@ -29,6 +30,7 @@ use crate::sabre_swap::swap_map::SwapMap;
 use crate::sabre_swap::{build_swap_map_inner, Heuristic, NodeBlockResults, SabreResult};
 
 #[pyfunction]
+#[pyo3(signature = (dag, neighbor_table, distance_matrix, heuristic, max_iterations, num_swap_trials, num_random_trials, seed=None, partial_layouts=vec![]))]
 pub fn sabre_layout_and_routing(
     py: Python,
     dag: &SabreDAG,
@@ -37,20 +39,24 @@ pub fn sabre_layout_and_routing(
     heuristic: &Heuristic,
     max_iterations: usize,
     num_swap_trials: usize,
-    num_layout_trials: usize,
+    num_random_trials: usize,
     seed: Option<u64>,
+    mut partial_layouts: Vec<Vec<Option<usize>>>,
 ) -> ([NLayout; 2], (SwapMap, PyObject, NodeBlockResults)) {
     let run_in_parallel = getenv_use_multiple_threads();
+    let mut starting_layouts: Vec<Vec<Option<usize>>> =
+        (0..num_random_trials).map(|_| vec![]).collect();
+    starting_layouts.append(&mut partial_layouts);
     let outer_rng = match seed {
         Some(seed) => Pcg64Mcg::seed_from_u64(seed),
         None => Pcg64Mcg::from_entropy(),
     };
     let seed_vec: Vec<u64> = outer_rng
         .sample_iter(&rand::distributions::Standard)
-        .take(num_layout_trials)
+        .take(starting_layouts.len())
         .collect();
     let dist = distance_matrix.as_array();
-    let res = if run_in_parallel && num_layout_trials > 1 {
+    let res = if run_in_parallel && starting_layouts.len() > 1 {
         seed_vec
             .into_par_iter()
             .enumerate()
@@ -66,6 +72,7 @@ pub fn sabre_layout_and_routing(
                         max_iterations,
                         num_swap_trials,
                         run_in_parallel,
+                        &starting_layouts[index],
                     ),
                 )
             })
@@ -80,7 +87,8 @@ pub fn sabre_layout_and_routing(
     } else {
         seed_vec
             .into_iter()
-            .map(|seed_trial| {
+            .enumerate()
+            .map(|(index, seed_trial)| {
                 layout_trial(
                     dag,
                     neighbor_table,
@@ -90,6 +98,7 @@ pub fn sabre_layout_and_routing(
                     max_iterations,
                     num_swap_trials,
                     run_in_parallel,
+                    &starting_layouts[index],
                 )
             })
             .min_by_key(|(_, result)| result.map.map.values().map(|x| x.len()).sum::<usize>())
@@ -114,12 +123,33 @@ fn layout_trial(
     max_iterations: usize,
     num_swap_trials: usize,
     run_swap_in_parallel: bool,
+    starting_layout: &[Option<usize>],
 ) -> ([NLayout; 2], SabreResult) {
     // Pick a random initial layout and fully populate ancillas in that layout too
     let num_physical_qubits = distance_matrix.shape()[0];
     let mut rng = Pcg64Mcg::seed_from_u64(seed);
-    let mut physical_qubits: Vec<usize> = (0..num_physical_qubits).collect();
-    physical_qubits.shuffle(&mut rng);
+    let physical_qubits = if !starting_layout.is_empty() {
+        let used_bits: HashSet<usize> = starting_layout
+            .iter()
+            .filter_map(|x| x.as_ref())
+            .copied()
+            .collect();
+        let mut free_bits: Vec<usize> = (0..num_physical_qubits)
+            .filter(|x| !used_bits.contains(x))
+            .collect();
+        free_bits.shuffle(&mut rng);
+        (0..num_physical_qubits)
+            .map(|x| match starting_layout.get(x) {
+                Some(phys) => phys.unwrap_or_else(|| free_bits.pop().unwrap()),
+                None => free_bits.pop().unwrap(),
+            })
+            .collect()
+    } else {
+        let mut physical_qubits: Vec<usize> = (0..num_physical_qubits).collect();
+        physical_qubits.shuffle(&mut rng);
+        physical_qubits
+    };
+
     let mut initial_layout = NLayout::from_logical_to_physical(physical_qubits);
     let new_dag_fn = |nodes| {
         // Because the current implementation of Sabre swap doesn't permute

qiskit/transpiler/passes/layout/sabre_layout.py

@@ -15,6 +15,8 @@
 
 import copy
 import logging
+import functools
+
 import numpy as np
 import rustworkx as rx
 
@@ -69,6 +71,34 @@ class SabreLayout(TransformationPass):
     layout pass. When specified this will use the specified routing pass to select an
     initial layout only and will not run multiple seed trials.
 
+    In addition to starting with a random initial `Layout` the pass can also take in
+    an additional list of starting layouts which will be used for additional
+    trials. If the ``sabre_starting_layout`` is present in the property set
+    when this pass is run, that will be used for additional trials. There will still
+    be ``layout_trials`` of full random starting layouts run and the contents of
+    ``sabre_starting_layout`` will be run in addition to those. The output which results
+    in the lowest amount of swap gates (whether from the random trials or the property
+    set starting point) will be used. The value for this property set field should be a
+    list of :class:`.Layout` objects representing the starting layouts to use. If a
+    virtual qubit is missing from an :class:`.Layout` object in the list a random qubit
+    will be selected.
+
+    Property Set Fields Read
+    ------------------------
+
+    ``sabre_starting_layout`` (``list[Layout]``)
+        An optional list of :class:`~.Layout` objects to use for additional layout trials. This is
+        in addition to the full random trials specified with the ``layout_trials`` argument.
+
+    Property Set Values Written
+    ---------------------------
+
+    ``layout`` (:class:`.Layout`)
+        The chosen initial mapping of virtual to physical qubits, including the ancilla allocation.
+
+    ``final_layout`` (:class:`.Layout`)
+        A permutation of how swaps have been applied to the input qubits at the end of the circuit.
+
     **References:**
 
     [1] Li, Gushu, Yufei Ding, and Yuan Xie. "Tackling the qubit mapping problem
@@ -227,10 +257,16 @@ def run(self, dag):
             target.make_symmetric()
         else:
             target = self.coupling_map
+        inner_run = self._inner_run
+        if "sabre_starting_layout" in self.property_set:
+            inner_run = functools.partial(
+                self._inner_run, starting_layout=self.property_set["sabre_starting_layout"]
+            )
+
         layout_components = disjoint_utils.run_pass_over_connected_components(
             dag,
             target,
-            self._inner_run,
+            inner_run,
         )
         initial_layout_dict = {}
         final_layout_dict = {}
@@ -284,7 +320,7 @@ def run(self, dag):
         disjoint_utils.combine_barriers(mapped_dag, retain_uuid=False)
         return mapped_dag
 
-    def _inner_run(self, dag, coupling_map):
+    def _inner_run(self, dag, coupling_map, starting_layout=None):
         if not coupling_map.is_symmetric:
             # deepcopy is needed here to avoid modifications updating
             # shared references in passes which require directional
@@ -294,6 +330,21 @@ def _inner_run(self, dag, coupling_map):
         neighbor_table = NeighborTable(rx.adjacency_matrix(coupling_map.graph))
         dist_matrix = coupling_map.distance_matrix
         original_qubit_indices = {bit: index for index, bit in enumerate(dag.qubits)}
+        partial_layouts = []
+        if starting_layout is not None:
+            coupling_map_reverse_mapping = {
+                coupling_map.graph[x]: x for x in coupling_map.graph.node_indices()
+            }
+            for layout in starting_layout:
+                virtual_bits = layout.get_virtual_bits()
+                out_layout = [None] * len(dag.qubits)
+                for bit, phys in virtual_bits.items():
+                    pos = original_qubit_indices.get(bit, None)
+                    if pos is None:
+                        continue
+                    out_layout[pos] = coupling_map_reverse_mapping[phys]
+                partial_layouts.append(out_layout)
+
         sabre_dag, circuit_to_dag_dict = _build_sabre_dag(
             dag,
             coupling_map.size(),
@@ -308,6 +359,7 @@ def _inner_run(self, dag, coupling_map):
             self.swap_trials,
             self.layout_trials,
             self.seed,
+            partial_layouts,
         )
 
         # Apply initial layout selected.

releasenotes/notes/seed-sabre-with-layout-17d46e1a6f516b0e.yaml

@@ -0,0 +1,38 @@
+---
+features:
+  - |
+    Added support to the :class:`.SabreLayout` pass to add trials with specified
+    starting layouts. The :class:`.SabreLayout` transpiler pass typically
+    runs multiple layout trials that all start with fully random layouts which
+    then use a routing pass to permute that layout instead of inserting swaps
+    to find a layout which will result in fewer swap gates. This new feature
+    enables running an :class:`.AnalysisPass` prior to :class:`.SabreLayout`
+    which sets the ``"sabre_starting_layout"`` field in the property set
+    to provide the :class:`.SabreLayout` with additional starting layouts
+    to use in its internal trials. For example, if you wanted to run
+    :class:`.DenseLayout` as the starting point for one trial in
+    :class:`.SabreLayout` you would do something like::
+
+        from qiskit.providers.fake_provider import FakeSherbrooke
+        from qiskit.transpiler import AnalysisPass, PassManager
+        from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
+        from qiskit.transpiler.passes import DenseLayout
+
+        class SabreDenseLayoutTrial(AnalysisPass):
+
+          def __init__(self, target):
+              self.dense_pass = DenseLayout(target=target)
+              super().__init__()
+
+          def run(self, dag):
+              self.dense_pass.run(dag)
+              self.property_set["sabre_starting_layout"] = [self.dense_pass.property_set["layout"]]
+
+        backend = FakeSherbrooke()
+        opt_level_1 = generate_preset_pass_manager(1, backend)
+        pre_layout = PassManager([SabreDenseLayoutTrial(backend.target)])
+        opt_level_1.pre_layout = pre_layout
+
+    Then when the ``opt_level_1`` :class:`.StagedPassManager` is run with a circuit the output
+    of the :class:`.DenseLayout` pass will be used for one of the :class:`.SabreLayout` trials
+    in addition to the 5 fully random trials that run by default in optimization level 1.

test/python/transpiler/test_sabre_layout.py

@@ -15,8 +15,8 @@
 import unittest
 
 from qiskit import QuantumRegister, QuantumCircuit
-from qiskit.transpiler import CouplingMap
-from qiskit.transpiler.passes import SabreLayout
+from qiskit.transpiler import CouplingMap, AnalysisPass, PassManager
+from qiskit.transpiler.passes import SabreLayout, DenseLayout
 from qiskit.transpiler.exceptions import TranspilerError
 from qiskit.converters import circuit_to_dag
 from qiskit.test import QiskitTestCase
@@ -94,6 +94,41 @@ def test_6q_circuit_20q_coupling(self):
         layout = pass_.property_set["layout"]
         self.assertEqual([layout[q] for q in circuit.qubits], [7, 8, 12, 6, 11, 13])
 
+    def test_6q_circuit_20q_coupling_with_partial(self):
+        """Test finds layout for 6q circuit on 20q device."""
+        #       ┌───┐┌───┐┌───┐┌───┐┌───┐
+        # q0_0: ┤ X ├┤ X ├┤ X ├┤ X ├┤ X ├
+        #       └─┬─┘└─┬─┘└─┬─┘└─┬─┘└─┬─┘
+        # q0_1: ──┼────■────┼────┼────┼──
+        #         │  ┌───┐  │    │    │
+        # q0_2: ──┼──┤ X ├──┼────■────┼──
+        #         │  └───┘  │         │
+        # q1_0: ──■─────────┼─────────┼──
+        #            ┌───┐  │         │
+        # q1_1: ─────┤ X ├──┼─────────■──
+        #            └───┘  │
+        # q1_2: ────────────■────────────
+        qr0 = QuantumRegister(3, "q0")
+        qr1 = QuantumRegister(3, "q1")
+        circuit = QuantumCircuit(qr0, qr1)
+        circuit.cx(qr1[0], qr0[0])
+        circuit.cx(qr0[1], qr0[0])
+        circuit.cx(qr1[2], qr0[0])
+        circuit.x(qr0[2])
+        circuit.cx(qr0[2], qr0[0])
+        circuit.x(qr1[1])
+        circuit.cx(qr1[1], qr0[0])
+
+        pm = PassManager(
+            [
+                DensePartialSabreTrial(CouplingMap(self.cmap20)),
+                SabreLayout(CouplingMap(self.cmap20), seed=0, swap_trials=32, layout_trials=32),
+            ]
+        )
+        pm.run(circuit)
+        layout = pm.property_set["layout"]
+        self.assertEqual([layout[q] for q in circuit.qubits], [7, 8, 12, 6, 11, 13])
+
     def test_6q_circuit_20q_coupling_with_target(self):
         """Test finds layout for 6q circuit on 20q device."""
         #       ┌───┐┌───┐┌───┐┌───┐┌───┐
@@ -218,6 +253,18 @@ def test_layout_many_search_trials(self):
         )
 
 
+class DensePartialSabreTrial(AnalysisPass):
+    """Pass to run dense layout as a sabre trial."""
+
+    def __init__(self, cmap):
+        self.dense_pass = DenseLayout(cmap)
+        super().__init__()
+
+    def run(self, dag):
+        self.dense_pass.run(dag)
+        self.property_set["sabre_starting_layout"] = [self.dense_pass.property_set["layout"]]
+
+
 class TestDisjointDeviceSabreLayout(QiskitTestCase):
     """Test SabreLayout with a disjoint coupling map."""
 
@@ -319,6 +366,28 @@ def test_too_large_components(self):
         with self.assertRaises(TranspilerError):
             layout_routing_pass(qc)
 
+    def test_with_partial_layout(self):
+        """Test a partial layout with a disjoint connectivity graph."""
+        qc = QuantumCircuit(8, name="double dhz")
+        qc.h(0)
+        qc.cz(0, 1)
+        qc.cz(0, 2)
+        qc.h(3)
+        qc.cx(3, 4)
+        qc.cx(3, 5)
+        qc.cx(3, 6)
+        qc.cx(3, 7)
+        qc.measure_all()
+        pm = PassManager(
+            [
+                DensePartialSabreTrial(self.dual_grid_cmap),
+                SabreLayout(self.dual_grid_cmap, seed=123456, swap_trials=1, layout_trials=1),
+            ]
+        )
+        pm.run(qc)
+        layout = pm.property_set["layout"]
+        self.assertEqual([layout[q] for q in qc.qubits], [3, 1, 2, 5, 4, 6, 7, 8])
+
 
 if __name__ == "__main__":
     unittest.main()