Oxidize SabreLayout pass

This commit modifies the SabreLayout pass when run without the
routing_pass argument to run primarily in Rust. This builds on top of
the rust version of SabreSwap previously added in Qiskit#7977, Qiskit#8388,
and Qiskit#8572. Internally, when the routing_pass argument is not set
SabreLayout will perform the full sabre algorithm both layout selection
and final swap mapping in rust and return the selected initial layout,
the final layout, the toplogical sorting used to traverse the circuit,
and a SwapMap for any swaps inserted. This is then used to build the
output circuit in place of running separate layout and routing passes.
The preset pass managers are updated to handle the new combined layout
and routing mode of operation for SabreLayout. The routing stage to the
preset pass managers remains intact, it will just operate as if a
perfect layout was selected and skip SabreSwap because the circuit is
already matching the connectivity constraints.

Besides just operating more quickly because the heavy lifting of the
algorithm operates more efficiently in a compiled language, doing this
in rust also lets change our parallelization model for running multiple
seed in Sabre. Just as in Qiskit#8572 we added support for SabreSwap to run
multiple parallel trials with different seeds this commit adds a
layout_trials argument to SabreLayout to try multiple seeds in parallel.
When this is used it parallelizes at the outer layer for each
layout/routing combination and the total minimal swap count seed is used.
So for example if you set swap_trials=5 and layout_trails=5 that will run
5 tasks in the threadpool with 5 different seeds for the outer layout run.
Inside that every time sabre swap is run (which will be multiple times
as part of layout plus the final routing run) it tries 5 different seeds
for each execution serially inside that parallel task. This should
hopefully further improve the quality of the transpiler output and better
match expectations for users who were previously calling transpile()
multiple times to emulate this behavior.

Implements Qiskit#9090

qiskit/__init__.py

@@ -37,6 +37,7 @@
 sys.modules["qiskit._accelerate.nlayout"] = qiskit._accelerate.nlayout
 sys.modules["qiskit._accelerate.stochastic_swap"] = qiskit._accelerate.stochastic_swap
 sys.modules["qiskit._accelerate.sabre_swap"] = qiskit._accelerate.sabre_swap
+sys.modules["qiskit._accelerate.sabre_layout"] = qiskit._accelerate.sabre_layout
 sys.modules["qiskit._accelerate.pauli_expval"] = qiskit._accelerate.pauli_expval
 sys.modules["qiskit._accelerate.dense_layout"] = qiskit._accelerate.dense_layout
 sys.modules["qiskit._accelerate.sparse_pauli_op"] = qiskit._accelerate.sparse_pauli_op

qiskit/transpiler/passes/layout/sabre_layout.py

@@ -15,22 +15,30 @@
 
 import logging
 import numpy as np
+import retworkx
 
 from qiskit.converters import dag_to_circuit
 from qiskit.transpiler.passes.layout.set_layout import SetLayout
 from qiskit.transpiler.passes.layout.full_ancilla_allocation import FullAncillaAllocation
 from qiskit.transpiler.passes.layout.enlarge_with_ancilla import EnlargeWithAncilla
 from qiskit.transpiler.passes.layout.apply_layout import ApplyLayout
-from qiskit.transpiler.passes.routing import SabreSwap
 from qiskit.transpiler.passmanager import PassManager
 from qiskit.transpiler.layout import Layout
-from qiskit.transpiler.basepasses import AnalysisPass
+from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.exceptions import TranspilerError
+from qiskit._accelerate.nlayout import NLayout
+from qiskit._accelerate.sabre_layout import sabre_layout_and_routing
+from qiskit._accelerate.sabre_swap import (
+    Heuristic,
+    NeighborTable,
+)
+from qiskit.transpiler.passes.routing.sabre_swap import process_swaps, apply_gate
+from qiskit.tools.parallel import CPU_COUNT
 
 logger = logging.getLogger(__name__)
 
 
-class SabreLayout(AnalysisPass):
+class SabreLayout(TransformationPass):
     """Choose a Layout via iterative bidirectional routing of the input circuit.
 
     Starting with a random initial `Layout`, the algorithm does a full routing
@@ -50,7 +58,13 @@ class SabreLayout(AnalysisPass):
     """
 
     def __init__(
-        self, coupling_map, routing_pass=None, seed=None, max_iterations=3, swap_trials=None
+        self,
+        coupling_map,
+        routing_pass=None,
+        seed=None,
+        max_iterations=3,
+        swap_trials=None,
+        layout_trials=None,
     ):
         """SabreLayout initializer.
 
@@ -71,27 +85,44 @@ def __init__(
                 on the number of trials run. This option is mutually exclusive
                 with the ``routing_pass`` argument and an error will be raised
                 if both are used.
+            layout_trials (int): The number of random seed trials to run
+                layout with.
 
         Raises:
             TranspilerError: If both ``routing_pass`` and ``swap_trials`` are
                 specified
         """
         super().__init__()
         self.coupling_map = coupling_map
-        if routing_pass is not None and swap_trials is not None:
+        self._neighbor_table = None
+        if self.coupling_map is not None:
+            self._neighbor_table = NeighborTable(retworkx.adjacency_matrix(self.coupling_map.graph))
+
+        if routing_pass is not None and (swap_trials is not None or layout_trials is not None):
             raise TranspilerError("Both routing_pass and swap_trials can't be set at the same time")
         self.routing_pass = routing_pass
         self.seed = seed
         self.max_iterations = max_iterations
         self.trials = swap_trials
-        self.swap_trials = swap_trials
+        if swap_trials is None:
+            self.swap_trials = CPU_COUNT
+        else:
+            self.swap_trials = swap_trials
+        if layout_trials is None:
+            self.layout_trials = CPU_COUNT
+        else:
+            self.layout_trials = layout_trials
 
     def run(self, dag):
         """Run the SabreLayout pass on `dag`.
 
         Args:
             dag (DAGCircuit): DAG to find layout for.
 
+        Returns:
+           DAGCircuit: The output dag if swap mapping was run
+            (otherwise the input dag is returned unmodified).
+
         Raises:
             TranspilerError: if dag wider than self.coupling_map
         """
@@ -101,44 +132,119 @@ def run(self, dag):
         # Choose a random initial_layout.
         if self.seed is None:
             self.seed = np.random.randint(0, np.iinfo(np.int32).max)
-        rng = np.random.default_rng(self.seed)
 
-        physical_qubits = rng.choice(self.coupling_map.size(), len(dag.qubits), replace=False)
-        physical_qubits = rng.permutation(physical_qubits)
-        initial_layout = Layout({q: dag.qubits[i] for i, q in enumerate(physical_qubits)})
+        if self.routing_pass is not None:
+            rng = np.random.default_rng(self.seed)
+
+            physical_qubits = rng.choice(self.coupling_map.size(), len(dag.qubits), replace=False)
+            physical_qubits = rng.permutation(physical_qubits)
+            initial_layout = Layout({q: dag.qubits[i] for i, q in enumerate(physical_qubits)})
 
-        if self.routing_pass is None:
-            self.routing_pass = SabreSwap(
-                self.coupling_map, "decay", seed=self.seed, fake_run=True, trials=self.swap_trials
-            )
-        else:
             self.routing_pass.fake_run = True
 
-        # Do forward-backward iterations.
-        circ = dag_to_circuit(dag)
-        rev_circ = circ.reverse_ops()
-        for _ in range(self.max_iterations):
-            for _ in ("forward", "backward"):
-                pm = self._layout_and_route_passmanager(initial_layout)
-                new_circ = pm.run(circ)
-
-                # Update initial layout and reverse the unmapped circuit.
-                pass_final_layout = pm.property_set["final_layout"]
-                final_layout = self._compose_layouts(
-                    initial_layout, pass_final_layout, new_circ.qregs
-                )
-                initial_layout = final_layout
-                circ, rev_circ = rev_circ, circ
+            # Do forward-backward iterations.
+            circ = dag_to_circuit(dag)
+            rev_circ = circ.reverse_ops()
+            for _ in range(self.max_iterations):
+                for _ in ("forward", "backward"):
+                    pm = self._layout_and_route_passmanager(initial_layout)
+                    new_circ = pm.run(circ)
 
-            # Diagnostics
-            logger.info("new initial layout")
-            logger.info(initial_layout)
+                    # Update initial layout and reverse the unmapped circuit.
+                    pass_final_layout = pm.property_set["final_layout"]
+                    final_layout = self._compose_layouts(
+                        initial_layout, pass_final_layout, new_circ.qregs
+                    )
+                    initial_layout = final_layout
+                    circ, rev_circ = rev_circ, circ
 
-        for qreg in dag.qregs.values():
-            initial_layout.add_register(qreg)
+                # Diagnostics
+                logger.info("new initial layout")
+                logger.info(initial_layout)
 
-        self.property_set["layout"] = initial_layout
-        self.routing_pass.fake_run = False
+            for qreg in dag.qregs.values():
+                initial_layout.add_register(qreg)
+            self.property_set["layout"] = initial_layout
+            self.routing_pass.fake_run = False
+            return dag
+        else:
+            dist_matrix = self.coupling_map.distance_matrix
+            original_qubit_indices = {bit: index for index, bit in enumerate(dag.qubits)}
+            original_clbit_indices = {bit: index for index, bit in enumerate(dag.clbits)}
+
+            dag_list = []
+            for node in dag.topological_op_nodes():
+                cargs = {original_clbit_indices[x] for x in node.cargs}
+                if node.op.condition is not None:
+                    for clbit in dag._bits_in_condition(node.op.condition):
+                        cargs.add(original_clbit_indices[clbit])
+
+                dag_list.append(
+                    (
+                        node._node_id,
+                        [original_qubit_indices[x] for x in node.qargs],
+                        cargs,
+                    )
+                )
+            ((initial_layout, final_layout), swap_map, gate_order) = sabre_layout_and_routing(
+                len(dag.clbits),
+                dag_list,
+                self._neighbor_table,
+                dist_matrix,
+                Heuristic.Decay,
+                self.seed,
+                self.max_iterations,
+                self.swap_trials,
+                self.layout_trials,
+            )
+            # Apply initial layout selected.
+            # this is a pseudo-pass manager to avoid the repeated round trip between
+            # dag and circuit and just use a dag
+            original_dag = dag
+            layout_dict = {}
+            num_qubits = len(dag.qubits)
+            for k, v in initial_layout.layout_mapping():
+                if k < num_qubits:
+                    layout_dict[dag.qubits[k]] = v
+            initital_layout = Layout(layout_dict)
+            self.property_set["layout"] = initital_layout
+            ancilla_pass = FullAncillaAllocation(self.coupling_map)
+            ancilla_pass.property_set = self.property_set
+            dag = ancilla_pass.run(dag)
+            enlarge_pass = EnlargeWithAncilla()
+            enlarge_pass.property_set = ancilla_pass.property_set
+            dag = enlarge_pass.run(dag)
+            apply_pass = ApplyLayout()
+            apply_pass.property_set = enlarge_pass.property_set
+            dag = apply_pass.run(dag)
+            # Apply sabre swap ontop of circuit with sabre layout
+            final_layout_mapping = final_layout.layout_mapping()
+            self.property_set["final_layout"] = Layout(
+                {dag.qubits[k]: v for (k, v) in final_layout_mapping}
+            )
+            mapped_dag = dag.copy_empty_like()
+            canonical_register = dag.qregs["q"]
+            current_layout = Layout.generate_trivial_layout(canonical_register)
+            qubit_indices = {bit: idx for idx, bit in enumerate(canonical_register)}
+            layout_mapping = {
+                qubit_indices[k]: v for k, v in current_layout.get_virtual_bits().items()
+            }
+            original_layout = NLayout(layout_mapping, len(dag.qubits), self.coupling_map.size())
+            for node_id in gate_order:
+                node = original_dag._multi_graph[node_id]
+                process_swaps(
+                    swap_map,
+                    node,
+                    mapped_dag,
+                    original_layout,
+                    canonical_register,
+                    False,
+                    qubit_indices,
+                )
+                apply_gate(
+                    mapped_dag, node, original_layout, canonical_register, False, layout_dict
+                )
+            return mapped_dag
 
     def _layout_and_route_passmanager(self, initial_layout):
         """Return a passmanager for a full layout and routing.

qiskit/transpiler/passes/routing/sabre_swap.py

@@ -242,33 +242,63 @@ def run(self, dag):
         if not self.fake_run:
             for node_id in gate_order:
                 node = dag._multi_graph[node_id]
-                self._process_swaps(swap_map, node, mapped_dag, original_layout, canonical_register)
-                self._apply_gate(mapped_dag, node, original_layout, canonical_register)
-            return mapped_dag
-        return dag
-
-    def _process_swaps(self, swap_map, node, mapped_dag, current_layout, canonical_register):
-        if node._node_id in swap_map:
-            for swap in swap_map[node._node_id]:
-                swap_qargs = [canonical_register[swap[0]], canonical_register[swap[1]]]
-                self._apply_gate(
+                process_swaps(
+                    swap_map,
+                    node,
                     mapped_dag,
-                    DAGOpNode(op=SwapGate(), qargs=swap_qargs),
-                    current_layout,
+                    original_layout,
                     canonical_register,
+                    self.fake_run,
+                    self._qubit_indices,
                 )
-                current_layout.swap_logical(*swap)
-
-    def _apply_gate(self, mapped_dag, node, current_layout, canonical_register):
-        new_node = self._transform_gate_for_layout(node, current_layout, canonical_register)
-        if self.fake_run:
-            return new_node
-        return mapped_dag.apply_operation_back(new_node.op, new_node.qargs, new_node.cargs)
-
-    def _transform_gate_for_layout(self, op_node, layout, device_qreg):
-        """Return node implementing a virtual op on given layout."""
-        mapped_op_node = copy(op_node)
-        mapped_op_node.qargs = tuple(
-            device_qreg[layout.logical_to_physical(self._qubit_indices[x])] for x in op_node.qargs
-        )
-        return mapped_op_node
+                apply_gate(
+                    mapped_dag,
+                    node,
+                    original_layout,
+                    canonical_register,
+                    self.fake_run,
+                    self._qubit_indices,
+                )
+            return mapped_dag
+        return dag
+
+
+def process_swaps(
+    swap_map,
+    node,
+    mapped_dag,
+    current_layout,
+    canonical_register,
+    fake_run,
+    qubit_indices,
+):
+    """Process swaps from SwapMap."""
+    if node._node_id in swap_map:
+        for swap in swap_map[node._node_id]:
+            swap_qargs = [canonical_register[swap[0]], canonical_register[swap[1]]]
+            apply_gate(
+                mapped_dag,
+                DAGOpNode(op=SwapGate(), qargs=swap_qargs),
+                current_layout,
+                canonical_register,
+                fake_run,
+                qubit_indices,
+            )
+            current_layout.swap_logical(*swap)
+
+
+def apply_gate(mapped_dag, node, current_layout, canonical_register, fake_run, qubit_indices):
+    """Apply gate given the current layout."""
+    new_node = transform_gate_for_layout(node, current_layout, canonical_register, qubit_indices)
+    if fake_run:
+        return new_node
+    return mapped_dag.apply_operation_back(new_node.op, new_node.qargs, new_node.cargs)
+
+
+def transform_gate_for_layout(op_node, layout, device_qreg, qubit_indices):
+    """Return node implementing a virtual op on given layout."""
+    mapped_op_node = copy(op_node)
+    mapped_op_node.qargs = tuple(
+        device_qreg[layout.logical_to_physical(qubit_indices[x])] for x in op_node.qargs
+    )
+    return mapped_op_node

qiskit/transpiler/preset_passmanagers/level0.py

@@ -109,10 +109,17 @@ def _choose_layout_condition(property_set):
                 "layout", layout_method, pass_manager_config, optimization_level=0
             )
         else:
+
+            def _swap_mapped(property_set):
+                return property_set["final_layout"] is None
+
             layout = PassManager()
             layout.append(_given_layout)
             layout.append(_choose_layout, condition=_choose_layout_condition)
-            layout += common.generate_embed_passmanager(coupling_map)
+            embed = common.generate_embed_passmanager(coupling_map)
+            layout.append(
+                [pass_ for x in embed.passes() for pass_ in x["passes"]], condition=_swap_mapped
+            )
         routing = routing_pm
     else:
         layout = None

qiskit/transpiler/preset_passmanagers/level1.py

@@ -196,12 +196,19 @@ def _opt_control(property_set):
                 "layout", layout_method, pass_manager_config, optimization_level=1
             )
         else:
+
+            def _swap_mapped(property_set):
+                return property_set["final_layout"] is None
+
             layout = PassManager()
             layout.append(_given_layout)
             layout.append(_choose_layout_0, condition=_choose_layout_condition)
             layout.append(_choose_layout_1, condition=_layout_not_perfect)
             layout.append(_improve_layout, condition=_vf2_match_not_found)
-            layout += common.generate_embed_passmanager(coupling_map)
+            embed = common.generate_embed_passmanager(coupling_map)
+            layout.append(
+                [pass_ for x in embed.passes() for pass_ in x["passes"]], condition=_swap_mapped
+            )
 
         routing = routing_pm