Add option to SabreLayout to specify starting layouts (#10721)

* Add option to SabreLayout to specify starting layouts

The SabreLayout pass typically starts with each layout trial with a
random starting layout. However, in some cases starting with a specific
starting layout can result in better output quality than starting with
a fully random layout. To use this feature an analysis pass can set a
new `sabre_starting_layout` field in the property set before
`SabreLayout` with a list of layout objects that will add additional
trials using each layout as a starting layout.

* Fix lint

* Combine random layout and partial layouts

* Expand property set docs

Co-authored-by: Jake Lishman <[email protected]>

* Simplify branching logic for parallelism

* Use a slice for starting_layout in a trial

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

Co-authored-by: Jake Lishman <[email protected]>

* Make starting layout variable names plural

The starting layout is a list and can have more than one entry. To make
this clearer, this commit renames starting_layout -> starting_layouts
and sabre_starting_layout -> sabre_starting_layouts.

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

Co-authored-by: Jake Lishman <[email protected]>

* Update property set key in test

* Update test to only use partial layout

---------

Co-authored-by: Jake Lishman <[email protected]>
Co-authored-by: Jake Lishman <[email protected]>

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, PyArray, PyReadonlyArray2};
 use pyo3::prelude::*;
@@ -28,6 +29,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,
@@ -36,20 +38,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<u32>>>,
 ) -> (NLayout, PyObject, (SwapMap, PyObject, NodeBlockResults)) {
     let run_in_parallel = getenv_use_multiple_threads();
+    let mut starting_layouts: Vec<Vec<Option<u32>>> =
+        (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()
@@ -65,6 +71,7 @@ pub fn sabre_layout_and_routing(
                         max_iterations,
                         num_swap_trials,
                         run_in_parallel,
+                        &starting_layouts[index],
                     ),
                 )
             })
@@ -79,7 +86,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,
@@ -89,6 +97,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,15 +123,38 @@ fn layout_trial(
     max_iterations: usize,
     num_swap_trials: usize,
     run_swap_in_parallel: bool,
+    starting_layout: &[Option<u32>],
 ) -> (NLayout, Vec<PhysicalQubit>, SabreResult) {
     let num_physical_qubits: u32 = distance_matrix.shape()[0].try_into().unwrap();
     let mut rng = Pcg64Mcg::seed_from_u64(seed);
 
     // Pick a random initial layout including a full ancilla allocation.
     let mut initial_layout = {
-        let mut physical_qubits: Vec<PhysicalQubit> =
-            (0..num_physical_qubits).map(PhysicalQubit::new).collect();
-        physical_qubits.shuffle(&mut rng);
+        let physical_qubits: Vec<PhysicalQubit> = if !starting_layout.is_empty() {
+            let used_bits: HashSet<u32> = starting_layout
+                .iter()
+                .filter_map(|x| x.as_ref())
+                .copied()
+                .collect();
+            let mut free_bits: Vec<u32> = (0..num_physical_qubits)
+                .filter(|x| !used_bits.contains(x))
+                .collect();
+            free_bits.shuffle(&mut rng);
+            (0..num_physical_qubits)
+                .map(|x| {
+                    let bit_index = match starting_layout.get(x as usize) {
+                        Some(phys) => phys.unwrap_or_else(|| free_bits.pop().unwrap()),
+                        None => free_bits.pop().unwrap(),
+                    };
+                    PhysicalQubit::new(bit_index)
+                })
+                .collect()
+        } else {
+            let mut physical_qubits: Vec<PhysicalQubit> =
+                (0..num_physical_qubits).map(PhysicalQubit::new).collect();
+            physical_qubits.shuffle(&mut rng);
+            physical_qubits
+        };
         NLayout::from_virtual_to_physical(physical_qubits).unwrap()
     };
 

qiskit/transpiler/passes/layout/sabre_layout.py

@@ -16,6 +16,7 @@
 import copy
 import dataclasses
 import logging
+import functools
 import time
 
 import numpy as np
@@ -74,6 +75,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_layouts`` 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_layouts`` 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_layouts`` (``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
@@ -232,8 +261,12 @@ def run(self, dag):
             target.make_symmetric()
         else:
             target = self.coupling_map
-
-        components = disjoint_utils.run_pass_over_connected_components(dag, target, self._inner_run)
+        inner_run = self._inner_run
+        if "sabre_starting_layouts" in self.property_set:
+            inner_run = functools.partial(
+                self._inner_run, starting_layouts=self.property_set["sabre_starting_layouts"]
+            )
+        components = disjoint_utils.run_pass_over_connected_components(dag, target, inner_run)
         self.property_set["layout"] = Layout(
             {
                 component.dag.qubits[logic]: component.coupling_map.graph[phys]
@@ -314,7 +347,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_layouts=None):
         if not coupling_map.is_symmetric:
             # deepcopy is needed here to avoid modifications updating
             # shared references in passes which require directional
@@ -323,8 +356,26 @@ def _inner_run(self, dag, coupling_map):
             coupling_map.make_symmetric()
         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_layouts is not None:
+            coupling_map_reverse_mapping = {
+                coupling_map.graph[x]: x for x in coupling_map.graph.node_indices()
+            }
+            for layout in starting_layouts:
+                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(), {bit: index for index, bit in enumerate(dag.qubits)}
+            dag,
+            coupling_map.size(),
+            original_qubit_indices,
         )
         sabre_start = time.perf_counter()
         (initial_layout, final_permutation, sabre_result) = sabre_layout_and_routing(
@@ -336,6 +387,7 @@ def _inner_run(self, dag, coupling_map):
             self.swap_trials,
             self.layout_trials,
             self.seed,
+            partial_layouts,
         )
         sabre_stop = time.perf_counter()
         logger.debug(

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_layouts"] = [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=0),
+            ]
+        )
+        pm.run(circuit)
+        layout = pm.property_set["layout"]
+        self.assertEqual([layout[q] for q in circuit.qubits], [1, 3, 5, 2, 6, 0])
+
     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_layouts"] = [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()